diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index fe69a3ace2..285a06ce31 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -16,13 +16,16 @@ jobs: name: 'Core' runs-on: ubuntu-latest container: - image: nvcr.io/nvidia/cuda:12.1.0-devel-ubuntu22.04 + image: nvcr.io/nvidia/cuda:12.9.2-devel-ubuntu22.04 options: --user root steps: - name: 'Dependencies' run: | apt-get update - apt-get install -y git python3.9 pip cudnn9-cuda-12 + apt-mark unhold libnccl2 libnccl-dev + apt-get install -y git python3 python3-pip cudnn9-cuda-12 \ + libnccl2=2.30.7-1+cuda12.9 \ + libnccl-dev=2.30.7-1+cuda12.9 pip install cmake==3.21.0 pybind11[global] ninja "nvidia-cudnn-frontend>=1.25.0" - name: 'Checkout' uses: actions/checkout@v3 @@ -70,13 +73,16 @@ jobs: - name: Start named container run: | - docker run -v $(pwd):$(pwd) -w $(pwd) --name builder -d nvcr.io/nvidia/cuda:12.8.0-devel-ubuntu22.04 sleep infinity + docker run -v $(pwd):$(pwd) -w $(pwd) --name builder -d nvcr.io/nvidia/cuda:12.9.2-devel-ubuntu22.04 sleep infinity - name: 'Dependencies' run: | docker exec builder bash -c '\ apt-get update && \ - apt-get install -y git python3.9 pip cudnn9-cuda-12 && \ + apt-mark unhold libnccl2 libnccl-dev && \ + apt-get install -y git python3 python3-pip cudnn9-cuda-12 \ + libnccl2=2.30.7-1+cuda12.9 \ + libnccl-dev=2.30.7-1+cuda12.9 && \ pip install cmake torch ninja pydantic importlib-metadata>=1.0 packaging pybind11 numpy einops onnxscript "nvidia-cudnn-frontend>=1.25.0" && \ apt-get clean \ ' @@ -92,7 +98,7 @@ jobs: name: 'JAX' runs-on: ubuntu-latest container: - image: ghcr.io/nvidia/jax:jax + image: ghcr.io/nvidia/jax:jax-2026-07-21 options: --user root steps: - name: 'Dependencies' @@ -143,7 +149,7 @@ jobs: - name: Start named container run: | - docker run -v $(pwd):$(pwd) -w $(pwd) --name builder -d ghcr.io/nvidia/jax:jax sleep infinity + docker run -v $(pwd):$(pwd) -w $(pwd) --name builder -d ghcr.io/nvidia/jax:jax-2026-07-21 sleep infinity - name: 'Dependencies' run: | diff --git a/.github/workflows/trigger-ci.yml b/.github/workflows/trigger-ci.yml index 68d1d7d71f..f0cb431425 100644 --- a/.github/workflows/trigger-ci.yml +++ b/.github/workflows/trigger-ci.yml @@ -62,6 +62,8 @@ jobs: || github.actor == 'jomitchellnv' || github.actor == 'fheinecke' || github.actor == 'janekb04' + || github.actor == 'YangFei1990' + || github.actor == 'sraman-rgb' ) steps: - name: Check if comment is issued by authorized person diff --git a/.gitmodules b/.gitmodules index d01719cc63..570dd1ce09 100644 --- a/.gitmodules +++ b/.gitmodules @@ -30,6 +30,6 @@ [submodule "3rdparty/ck_jit"] path = 3rdparty/ck_jit url = https://github.com/ROCm/ck-jit.git -[submodule "3rdparty/nccl"] - path = 3rdparty/nccl - url = https://github.com/NVIDIA/nccl.git +[submodule "3rdparty/nccl-extensions"] + path = 3rdparty/nccl-extensions + url = https://github.com/NVIDIA/nccl-extensions.git diff --git a/3rdparty/nccl b/3rdparty/nccl deleted file mode 160000 index b87848fbc5..0000000000 --- a/3rdparty/nccl +++ /dev/null @@ -1 +0,0 @@ -Subproject commit b87848fbc52da65b5a898b4ac6633fcf51cec4ed diff --git a/3rdparty/nccl-extensions b/3rdparty/nccl-extensions new file mode 160000 index 0000000000..9f47d6eb3b --- /dev/null +++ b/3rdparty/nccl-extensions @@ -0,0 +1 @@ +Subproject commit 9f47d6eb3b60962d8157a579b4caaaa4ae6b19f4 diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst index 14f1ee08d2..8e2ae5574d 100644 --- a/CONTRIBUTING.rst +++ b/CONTRIBUTING.rst @@ -13,7 +13,7 @@ Coding Guidelines ----------------- * We follow `Google C++ Style Guide `_. When no - rules can be found, follow the already occuring conventions. If there is no precedence in our + rules can be found, follow the already occurring conventions. If there is no precedence in our codebase we are open to discussion. * Prior to your contribution, please make sure that the code passes the linter check. We do both C++ and Python linting. To invoke the check, please use diff --git a/README.rst b/README.rst index 345d95515c..e94766413d 100644 --- a/README.rst +++ b/README.rst @@ -510,7 +510,7 @@ System Requirements * Compiler: GCC 9+ or Clang 10+ with C++17 support * Python: 3.12 recommended -* **Source Build Requirements:** CMake 3.18+, Ninja, Git 2.17+, pybind11 2.6.0+ +* **Source Build Requirements:** CMake 3.18+, Ninja, Git 2.17+, pybind11 2.6.0+, nvidia-cudnn-frontend 1.25.0+ * **Notes:** FP8 features require Compute Capability 8.9+ (Ada/Hopper/Blackwell) diff --git a/benchmarks/linear/benchmark_graph_safe_grouped_linear.py b/benchmarks/linear/benchmark_graph_safe_grouped_mlp.py similarity index 97% rename from benchmarks/linear/benchmark_graph_safe_grouped_linear.py rename to benchmarks/linear/benchmark_graph_safe_grouped_mlp.py index d8230c38fe..00f7f516c8 100644 --- a/benchmarks/linear/benchmark_graph_safe_grouped_linear.py +++ b/benchmarks/linear/benchmark_graph_safe_grouped_mlp.py @@ -13,21 +13,21 @@ Example: - python benchmarks/linear/benchmark_graph_safe_grouped_linear.py + python benchmarks/linear/benchmark_graph_safe_grouped_mlp.py Forward-only: - python benchmarks/linear/benchmark_graph_safe_grouped_linear.py --fwd-only + python benchmarks/linear/benchmark_graph_safe_grouped_mlp.py --fwd-only Nsight Systems: (optionally: unset DEBUGINFOD_URLS) nsys profile \ - --output=./benchmarks/linear/graph_safe_grouped_linear_mxfp8 \ + --output=./benchmarks/linear/graph_safe_grouped_mlp_mxfp8 \ --force-overwrite true \ --trace=cuda,nvtx,cudnn,cublas \ - python benchmarks/linear/benchmark_graph_safe_grouped_linear.py --profile + python benchmarks/linear/benchmark_graph_safe_grouped_mlp.py --profile """ # Match the Qwen MXFP8 SFT launch toggles before importing TE. diff --git a/benchmarks/linear/benchmark_grouped_gemm_kernels.py b/benchmarks/linear/benchmark_grouped_gemm_kernels.py new file mode 100644 index 0000000000..88b3f8a466 --- /dev/null +++ b/benchmarks/linear/benchmark_grouped_gemm_kernels.py @@ -0,0 +1,581 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Benchmark multi-stream and cuBLASLt grouped GEMM kernels. + +Unlike ``benchmark_grouped_linear_cublas_grouped_gemm.py``, this benchmark does +not construct a GroupedLinear module or invoke autograd. Inputs are allocated +and, for MXFP8, quantized and GEMM-swizzled before timing. The timed region only +calls TE's low-level grouped GEMM wrappers. + +The default shape is Qwen3.5-397B-A17B with sequence length 4096, top-10 +routing, and EP32. Each rank owns 16 experts and processes 2560 rows per expert. +All six expert GEMMs are measured independently: + +============== ====== ====== ====== ====== ================================== +Projection Pass Layout M K N +============== ====== ====== ====== ====== ================================== +FC1 fwd TN 2560 4096 2048 +FC1 dgrad NN 2560 4096 2048 +FC1 wgrad NT 2560 4096 2048 +FC2 fwd TN 2560 1024 4096 +FC2 dgrad NN 2560 1024 4096 +FC2 wgrad NT 2560 1024 4096 +============== ====== ====== ====== ====== ================================== + +The layouts describe TE's operands for each expert: + +* Forward (TN): A=``weight[N,K]``, B=``input[M,K]``; + ``input @ weight.T -> output[M,N]``. +* Dgrad (NN): A=``weight[N,K]``, B=``dy[M,N]``; + ``dy @ weight -> dx[M,K]``. +* Wgrad (NT): A=``input[M,K]``, B=``dy[M,N]``; + ``dy.T @ input -> dweight[N,K]``. + +The multi-stream path receives lists of discrete tensors. The cuBLASLt grouped +path receives packed GroupedTensor activations and discrete weights/wgrads, +matching GroupedLinear with discrete parameters. + +Examples +-------- +Run the full BF16 and MXFP8 matrix: + + python benchmarks/linear/benchmark_grouped_gemm_kernels.py + +Run only MXFP8 FC1 wgrad: + + python benchmarks/linear/benchmark_grouped_gemm_kernels.py \ + --precision mxfp8 --projection fc1 --gemm wgrad +""" + +import argparse +import os +from dataclasses import dataclass +from typing import Any, Optional + +import pandas as pd +import torch +import torch.utils.benchmark as benchmark + +from transformer_engine.common.recipe import MXFP8BlockScaling, Recipe +from transformer_engine.pytorch import MXFP8Quantizer +from transformer_engine.pytorch.cpp_extensions import ( + general_grouped_gemm, + general_grouped_gemm_for_grouped_tensor, +) +from transformer_engine.pytorch.module import is_module_grouped_tensor_path_supported +from transformer_engine.pytorch.module.base import ( + _2X_ACC_DGRAD, + _2X_ACC_FPROP, + _2X_ACC_WGRAD, +) +from transformer_engine.pytorch.quantization import FP8GlobalStateManager +from transformer_engine.pytorch.tensor import GroupedTensor, GroupedTensorStorage +import transformer_engine_torch as tex + + +QWEN_NUM_EXPERTS = 512 +QWEN_TOP_K = 10 +QWEN_SEQUENCE_LENGTH = 4096 +QWEN_EXPERT_PARALLEL_SIZE = 32 +QWEN_HIDDEN_SIZE = 4096 +QWEN_MOE_INTERMEDIATE_SIZE = 1024 + + +@dataclass(frozen=True) +class GemmSpec: + """One expert GEMM in the Qwen routed MLP.""" + + projection: str + gemm: str + layout: str + k: int + n: int + + +GEMM_SPECS = tuple( + GemmSpec(projection, gemm, layout, k, n) + for projection, k, n in ( + ("fc1", QWEN_HIDDEN_SIZE, 2 * QWEN_MOE_INTERMEDIATE_SIZE), + ("fc2", QWEN_MOE_INTERMEDIATE_SIZE, QWEN_HIDDEN_SIZE), + ) + for gemm, layout in (("fwd", "TN"), ("dgrad", "NN"), ("wgrad", "NT")) +) + + +@dataclass +class PreparedGemm: + """Preallocated operands and outputs for both grouped GEMM paths.""" + + spec: GemmSpec + split_sizes: list[int] + multistream_a: list[Any] + multistream_b: list[Any] + multistream_out: list[torch.Tensor] + grouped_a: Any + grouped_b: GroupedTensorStorage + grouped_out: Any + + +def _make_grouped_bf16( + packed: Optional[torch.Tensor], + split_sizes: list[int], + last_dim: int, +) -> GroupedTensorStorage: + """Create packed GroupedTensor storage and optionally initialize its data.""" + first_dims = torch.tensor(split_sizes, dtype=torch.int64, device="cuda") + grouped = GroupedTensor.make_grouped_tensor( + num_tensors=len(split_sizes), + first_dims=first_dims, + last_dims=None, + logical_first_dim=sum(split_sizes), + logical_last_dim=last_dim, + quantizer=None, + device=torch.device("cuda"), + dtype=torch.bfloat16, + ) + if packed is not None: + grouped.rowwise_data.view(-1).copy_(packed.reshape(-1)) + return grouped + + +def _new_mxfp8_quantizer(*, rowwise: bool, columnwise: bool) -> MXFP8Quantizer: + """Construct a quantizer that emits GEMM-ready scales before timing.""" + quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=rowwise, + columnwise=columnwise, + ) + quantizer.optimize_for_gemm = True + return quantizer + + +def _quantize_discrete_mxfp8( + tensors: list[torch.Tensor], + *, + rowwise: bool, + columnwise: bool, +) -> list[Any]: + """Quantize per-expert tensors outside the timed region.""" + quantizer = _new_mxfp8_quantizer(rowwise=rowwise, columnwise=columnwise) + return [quantizer(tensor) for tensor in tensors] + + +def _quantize_grouped_mxfp8( + packed: torch.Tensor, + split_sizes: list[int], + *, + rowwise: bool, + columnwise: bool, +) -> GroupedTensorStorage: + """Quantize one packed activation into GEMM-ready grouped storage.""" + quantizer = _new_mxfp8_quantizer(rowwise=rowwise, columnwise=columnwise) + first_dims = torch.tensor(split_sizes, dtype=torch.int64, device="cuda") + return tex.group_quantize( + packed, + quantizer, + len(split_sizes), + first_dims, + ) + + +def _make_high_precision_operands( + spec: GemmSpec, + split_sizes: list[int], +) -> tuple[list[torch.Tensor], Optional[torch.Tensor], torch.Tensor]: + """Create canonical BF16 operands shared by both precision paths.""" + num_experts = len(split_sizes) + total_rows = sum(split_sizes) + + if spec.gemm in ("fwd", "dgrad"): + weights = [ + torch.randn(spec.n, spec.k, dtype=torch.bfloat16, device="cuda") + for _ in range(num_experts) + ] + b_width = spec.k if spec.gemm == "fwd" else spec.n + packed_b = torch.randn(total_rows, b_width, dtype=torch.bfloat16, device="cuda") + return weights, None, packed_b + + packed_x = torch.randn(total_rows, spec.k, dtype=torch.bfloat16, device="cuda") + packed_dy = torch.randn(total_rows, spec.n, dtype=torch.bfloat16, device="cuda") + return [], packed_x, packed_dy + + +def _operand_usage(spec: GemmSpec) -> tuple[tuple[bool, bool], tuple[bool, bool]]: + """Return rowwise/columnwise MXFP8 storage used by A and B.""" + if spec.gemm == "fwd": + return (True, False), (True, False) + if spec.gemm == "dgrad": + return (False, True), (True, False) + return (False, True), (False, True) + + +def _allocate_outputs( + spec: GemmSpec, + split_sizes: list[int], +) -> tuple[list[torch.Tensor], Any]: + """Allocate output storage matching each low-level API's native contract.""" + num_experts = len(split_sizes) + if spec.gemm == "wgrad": + multistream_out = [ + torch.empty(spec.n, spec.k, dtype=torch.bfloat16, device="cuda") + for _ in range(num_experts) + ] + grouped_out = [ + torch.empty(spec.n, spec.k, dtype=torch.bfloat16, device="cuda") + for _ in range(num_experts) + ] + return multistream_out, grouped_out + + out_features = spec.n if spec.gemm == "fwd" else spec.k + multistream_out = [ + torch.empty(sum(split_sizes), out_features, dtype=torch.bfloat16, device="cuda") + ] + grouped_out = _make_grouped_bf16(None, split_sizes, out_features) + return multistream_out, grouped_out + + +def _prepare_gemm( + spec: GemmSpec, + split_sizes: list[int], + precision: str, +) -> PreparedGemm: + """Prepare one GEMM without leaving quantization inside the timed region.""" + hp_a, packed_a, packed_b = _make_high_precision_operands(spec, split_sizes) + a_usage, b_usage = _operand_usage(spec) + + if spec.gemm in ("fwd", "dgrad"): + if precision == "mxfp8": + multistream_a = _quantize_discrete_mxfp8( + hp_a, + rowwise=a_usage[0], + columnwise=a_usage[1], + ) + else: + multistream_a = hp_a + grouped_a = multistream_a + elif precision == "mxfp8": + assert packed_a is not None + grouped_a = _quantize_grouped_mxfp8( + packed_a, + split_sizes, + rowwise=a_usage[0], + columnwise=a_usage[1], + ) + else: + assert packed_a is not None + grouped_a = _make_grouped_bf16(packed_a, split_sizes, spec.k) + if spec.gemm == "wgrad": + multistream_a = list(grouped_a.split_into_quantized_tensors()) + + if precision == "mxfp8": + grouped_b = _quantize_grouped_mxfp8( + packed_b, + split_sizes, + rowwise=b_usage[0], + columnwise=b_usage[1], + ) + else: + b_width = spec.k if spec.gemm == "fwd" else spec.n + grouped_b = _make_grouped_bf16(packed_b, split_sizes, b_width) + # Give multi-stream zero-copy member views into the exact same packed operand used by + # cuBLASLt. This prevents quantization or input-data differences from entering the result. + multistream_b = list(grouped_b.split_into_quantized_tensors()) + + multistream_out, grouped_out = _allocate_outputs(spec, split_sizes) + prepared = PreparedGemm( + spec=spec, + split_sizes=split_sizes, + multistream_a=multistream_a, + multistream_b=multistream_b, + multistream_out=multistream_out, + grouped_a=grouped_a, + grouped_b=grouped_b, + grouped_out=grouped_out, + ) + if precision == "mxfp8": + _assert_mxfp8_operands_are_gemm_ready(prepared) + return prepared + + +def _iter_operands(operand: Any): + """Yield discrete members or one packed operand.""" + if isinstance(operand, (list, tuple)): + yield from operand + else: + yield operand + + +def _assert_mxfp8_operands_are_gemm_ready(prepared: PreparedGemm) -> None: + """Ensure scale swizzling cannot leak into the GEMM timing.""" + operands = ( + prepared.multistream_a, + prepared.multistream_b, + prepared.grouped_a, + prepared.grouped_b, + ) + for operand_group in operands: + for operand in _iter_operands(operand_group): + assert getattr( + operand, "_with_gemm_swizzled_scales", False + ), "MXFP8 operand was not prepared with GEMM-swizzled scales" + + +def _use_split_accumulator(spec: GemmSpec) -> bool: + """Match GroupedLinear's accumulator policy for each training GEMM.""" + if spec.gemm == "fwd": + return _2X_ACC_FPROP + if spec.gemm == "dgrad": + return _2X_ACC_DGRAD + return _2X_ACC_WGRAD + + +def _run_multistream(prepared: PreparedGemm, iterations: int) -> None: + """Launch the multi-stream grouped GEMM repeatedly.""" + spec = prepared.spec + for _ in range(iterations): + general_grouped_gemm( + prepared.multistream_a, + prepared.multistream_b, + prepared.multistream_out, + [None] * len(prepared.split_sizes), + torch.bfloat16, + layout=spec.layout, + m_splits=prepared.split_sizes, + grad=spec.gemm != "fwd", + single_output=spec.gemm != "wgrad", + use_split_accumulator=_use_split_accumulator(spec), + ) + + +def _run_cublas_grouped(prepared: PreparedGemm, iterations: int) -> None: + """Launch the device-described cuBLASLt grouped GEMM repeatedly.""" + for _ in range(iterations): + general_grouped_gemm_for_grouped_tensor( + prepared.grouped_a, + prepared.grouped_b, + prepared.grouped_out, + layout=prepared.spec.layout, + use_split_accumulator=_use_split_accumulator(prepared.spec), + ) + + +def _canonical_output(prepared: PreparedGemm, *, grouped: bool) -> torch.Tensor: + """Return either path's output in one comparable packed tensor.""" + output = prepared.grouped_out if grouped else prepared.multistream_out + if prepared.spec.gemm == "wgrad": + return torch.stack(output, dim=0) + if grouped: + out_features = prepared.spec.n if prepared.spec.gemm == "fwd" else prepared.spec.k + return output.rowwise_data.view(sum(prepared.split_sizes), out_features) + return output[0] + + +def _validate_outputs(prepared: PreparedGemm, precision: str) -> None: + """Check that execution-path selection does not change GEMM numerics unexpectedly.""" + _run_multistream(prepared, 1) + _run_cublas_grouped(prepared, 1) + if precision == "mxfp8": + tolerances = {"rtol": 0.125, "atol": 0.0675} + else: + tolerances = {"rtol": 1e-2, "atol": 1e-2} + torch.testing.assert_close( + _canonical_output(prepared, grouped=True), + _canonical_output(prepared, grouped=False), + **tolerances, + ) + + +def _benchmark_path( + prepared: PreparedGemm, + *, + path: str, + iterations_per_run: int, + warmup_iterations: int, + min_run_time: float, + profile: bool, + label: str, +) -> float: + """Return milliseconds for one grouped GEMM launch.""" + run = _run_multistream if path == "multistream" else _run_cublas_grouped + run(prepared, warmup_iterations) + torch.cuda.synchronize() + + if profile: + torch.cuda.nvtx.range_push(label) + timing = benchmark.Timer( + stmt="run(prepared, iterations_per_run)", + globals={ + "run": run, + "prepared": prepared, + "iterations_per_run": iterations_per_run, + }, + num_threads=1, + ).blocked_autorange(min_run_time=min_run_time) + if profile: + torch.cuda.nvtx.range_pop() + return timing.median * 1000 / iterations_per_run + + +def _gemm_tflops(spec: GemmSpec, total_rows: int, time_ms: float) -> float: + """Compute aggregate GEMM throughput across all local experts.""" + flops = 2 * total_rows * spec.k * spec.n + return flops / time_ms / 1e9 + + +def _shape_description(spec: GemmSpec, rows_per_expert: str) -> str: + """Format the actual per-expert operands passed to TE.""" + if spec.gemm == "fwd": + return ( + f"A[{spec.n},{spec.k}] B[{rows_per_expert},{spec.k}] -> D[{rows_per_expert},{spec.n}]" + ) + if spec.gemm == "dgrad": + return ( + f"A[{spec.n},{spec.k}] B[{rows_per_expert},{spec.n}] -> D[{rows_per_expert},{spec.k}]" + ) + return f"A[{rows_per_expert},{spec.k}] B[{rows_per_expert},{spec.n}] -> D[{spec.n},{spec.k}]" + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--precision", choices=("all", "bf16", "mxfp8"), default="all") + parser.add_argument("--projection", choices=("all", "fc1", "fc2"), default="all") + parser.add_argument("--gemm", choices=("all", "fwd", "dgrad", "wgrad"), default="all") + parser.add_argument("--iterations-per-run", type=int, default=1000) + parser.add_argument("--warmup-iterations", type=int, default=500) + parser.add_argument("--min-run-time", type=float, default=5.0) + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument( + "--m-splits", + type=str, + default=None, + help="Optional comma-separated per-expert rows; defaults to Qwen3.5-397B EP32.", + ) + parser.add_argument("--skip-correctness", action="store_true") + parser.add_argument("--profile", action="store_true", help="Add per-path NVTX ranges.") + parser.add_argument("--output-csv", type=str, default=None) + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for this benchmark.") + if os.getenv("NVTE_USE_CUTLASS_GROUPED_GEMM", "0") == "1": + raise RuntimeError( + "Unset NVTE_USE_CUTLASS_GROUPED_GEMM: this benchmark requires the " + "multi-stream cuBLAS baseline." + ) + if args.iterations_per_run < 1 or args.warmup_iterations < 1: + raise ValueError("iterations-per-run and warmup-iterations must be positive.") + + if args.m_splits is None: + num_local_experts = QWEN_NUM_EXPERTS // QWEN_EXPERT_PARALLEL_SIZE + rows_per_expert = QWEN_SEQUENCE_LENGTH * QWEN_TOP_K // num_local_experts + split_sizes = [rows_per_expert] * num_local_experts + else: + split_sizes = [int(value) for value in args.m_splits.split(",") if value] + if not split_sizes or any(value <= 0 for value in split_sizes): + raise ValueError("m_splits must contain positive integers.") + num_local_experts = len(split_sizes) + + total_rows = sum(split_sizes) + precisions = ("bf16", "mxfp8") if args.precision == "all" else (args.precision,) + specs = [ + spec + for spec in GEMM_SPECS + if (args.projection == "all" or spec.projection == args.projection) + and (args.gemm == "all" or spec.gemm == args.gemm) + ] + recipes: dict[str, Optional[Recipe]] = { + "bf16": None, + "mxfp8": MXFP8BlockScaling(), + } + mxfp8_available, reason_for_no_mxfp8 = FP8GlobalStateManager.is_mxfp8_available() + + uniform_rows = str(split_sizes[0]) if len(set(split_sizes)) == 1 else "variable" + print("Qwen3.5-397B-A17B grouped GEMM kernel benchmark") + print(f" local experts: {num_local_experts}") + print(f" total rows: {total_rows}") + print(f" m_splits: {split_sizes}") + print(" quantization is outside the timed region") + print() + for spec in specs: + print( + f" {spec.projection} {spec.gemm:5s} {spec.layout}: " + f"{_shape_description(spec, uniform_rows)}" + ) + print() + + rows = [] + for precision in precisions: + recipe = recipes[precision] + if precision == "mxfp8" and not mxfp8_available: + print(f"Skipping MXFP8: {reason_for_no_mxfp8}") + continue + if not is_module_grouped_tensor_path_supported(recipe, torch.bfloat16): + print( + f"Skipping {precision}: cuBLASLt grouped GEMM is unsupported on this " + "GPU or cuBLASLt version." + ) + continue + + for spec in specs: + torch.manual_seed(args.seed) + torch.cuda.manual_seed(args.seed) + prepared = _prepare_gemm(spec, split_sizes, precision) + if not args.skip_correctness: + _validate_outputs(prepared, precision) + + timings = {} + for path in ("multistream", "cublas_grouped_gemm"): + label = f"{precision}_{spec.projection}_{spec.gemm}_{path}" + timings[path] = _benchmark_path( + prepared, + path=path, + iterations_per_run=args.iterations_per_run, + warmup_iterations=args.warmup_iterations, + min_run_time=args.min_run_time, + profile=args.profile, + label=label, + ) + + speedup = timings["multistream"] / timings["cublas_grouped_gemm"] + for path, time_ms in timings.items(): + rows.append( + { + "precision": precision, + "projection": spec.projection, + "gemm": spec.gemm, + "layout": spec.layout, + "execution_path": path, + "num_local_experts": num_local_experts, + "total_rows": total_rows, + "k": spec.k, + "n": spec.n, + "time_ms": time_ms, + "tflops": _gemm_tflops(spec, total_rows, time_ms), + "speedup_vs_multistream": ( + speedup if path == "cublas_grouped_gemm" else 1.0 + ), + } + ) + + print( + f"{precision:6s} {spec.projection} {spec.gemm:5s}: " + f"multistream={timings['multistream']:.3f} ms, " + f"cuBLAS grouped={timings['cublas_grouped_gemm']:.3f} ms, " + f"speedup={speedup:.3f}x" + ) + + results = pd.DataFrame(rows) + print() + print(results.to_string(index=False)) + if args.output_csv is not None: + results.to_csv(args.output_csv, index=False) + print(f"\nWrote {args.output_csv}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/linear/benchmark_grouped_linear_cublas_grouped_gemm.py b/benchmarks/linear/benchmark_grouped_linear_cublas_grouped_gemm.py new file mode 100644 index 0000000000..1b400c0c24 --- /dev/null +++ b/benchmarks/linear/benchmark_grouped_linear_cublas_grouped_gemm.py @@ -0,0 +1,520 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Compare GroupedLinear's multi-stream and cuBLASLt grouped GEMM paths. + +The default problem is the routed-expert shape from Qwen3.5-397B-A17B with +sequence length 4096, top-10 routing, and expert parallelism 32: + + 512 global experts / EP32 = 16 local experts + 4096 tokens * top-10 / 16 local experts = 2560 rows per expert + +The 2560-row split is already 256-aligned, so both paths execute identical +GEMM work. FC1 and FC2 are benchmarked independently: + + FC1: 16 x (M=2560, K=4096, N=2048) + FC2: 16 x (M=2560, K=1024, N=4096) + +Both paths use discrete parameters initialized with identical values: + +* ``use_grouped_tensor=False`` receives CPU splits and launches the + multi-stream grouped GEMM implementation. +* ``use_grouped_tensor=True`` receives CUDA int64 splits and launches the + cuBLASLt grouped GEMM implementation. + +The MXFP8 cases use BF16 primary parameters. Each timing bundle refreshes the +MXFP8 weight cache on its first microbatch and reuses it on later microbatches. + +Examples +-------- +Run the complete BF16 and MXFP8 comparison: + + python benchmarks/linear/benchmark_grouped_linear_cublas_grouped_gemm.py + +Run only MXFP8 FC1 forward and backward: + + python benchmarks/linear/benchmark_grouped_linear_cublas_grouped_gemm.py \ + --precision mxfp8 --projection fc1 --mode fwd_bwd + +Profile one microbatch per timing invocation: + + nsys profile \ + --output=grouped_linear_cublas_grouped_gemm \ + --force-overwrite=true \ + --trace=cuda,nvtx,cublas \ + python benchmarks/linear/benchmark_grouped_linear_cublas_grouped_gemm.py \ + --precision mxfp8 --projection fc1 --mode fwd_bwd \ + --num-microbatches 1 --profile +""" + +import argparse +from contextlib import nullcontext +from dataclasses import dataclass +from typing import Optional + +import pandas as pd +import torch +import torch.utils.benchmark as benchmark + +import transformer_engine.pytorch as te +from transformer_engine.common.recipe import MXFP8BlockScaling, Recipe +from transformer_engine.pytorch.module import ( + GroupedLinear, + is_module_grouped_tensor_path_supported, +) +from transformer_engine.pytorch.quantization import FP8GlobalStateManager + + +QWEN_NUM_EXPERTS = 512 +QWEN_TOP_K = 10 +QWEN_SEQUENCE_LENGTH = 4096 +QWEN_EXPERT_PARALLEL_SIZE = 32 +QWEN_HIDDEN_SIZE = 4096 +QWEN_MOE_INTERMEDIATE_SIZE = 1024 + + +@dataclass(frozen=True) +class Projection: + """GroupedLinear dimensions for one Qwen routed-expert projection.""" + + name: str + in_features: int + out_features: int + + +PROJECTIONS = { + "fc1": Projection( + name="fc1", + in_features=QWEN_HIDDEN_SIZE, + out_features=2 * QWEN_MOE_INTERMEDIATE_SIZE, + ), + "fc2": Projection( + name="fc2", + in_features=QWEN_MOE_INTERMEDIATE_SIZE, + out_features=QWEN_HIDDEN_SIZE, + ), +} + + +@dataclass(frozen=True) +class ExecutionPath: + """GroupedLinear path and the corresponding m_splits representation.""" + + name: str + use_grouped_tensor: bool + + +EXECUTION_PATHS = ( + ExecutionPath(name="multistream", use_grouped_tensor=False), + ExecutionPath(name="cublas_grouped_gemm", use_grouped_tensor=True), +) + + +def _quantization_context(recipe: Optional[Recipe]): + """Construct a fresh quantization context for one invocation.""" + if recipe is None: + return nullcontext() + return te.autocast(enabled=True, recipe=recipe) + + +def _make_m_splits( + *, + use_grouped_tensor: bool, + split_sizes: list[int], +) -> torch.Tensor: + """Use the split representation consumed natively by each execution path.""" + device = "cuda" if use_grouped_tensor else "cpu" + return torch.tensor(split_sizes, dtype=torch.int64, device=device) + + +def _build_layer( + *, + projection: Projection, + num_local_experts: int, + use_grouped_tensor: bool, +) -> GroupedLinear: + """Construct a discrete-parameter GroupedLinear module.""" + return GroupedLinear( + num_local_experts, + projection.in_features, + projection.out_features, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + use_grouped_tensor=use_grouped_tensor, + ) + + +def _run_microbatches( + layer: GroupedLinear, + x: torch.Tensor, + m_splits: torch.Tensor, + grad_output: torch.Tensor, + *, + recipe: Optional[Recipe], + mode: str, + num_microbatches: int, +) -> torch.Tensor: + """Run one timed bundle while preserving realistic MXFP8 weight caching.""" + layer.zero_grad(set_to_none=True) + x.grad = None + + if mode == "fwd": + with torch.no_grad(), _quantization_context(recipe): + for microbatch in range(num_microbatches): + output = layer( + x, + m_splits, + is_first_microbatch=(microbatch == 0), + ) + return output + + with _quantization_context(recipe): + for microbatch in range(num_microbatches): + output = layer( + x, + m_splits, + is_first_microbatch=(microbatch == 0), + ) + output.backward(grad_output) + return output + + +def _run_correctness_step( + layer: GroupedLinear, + x: torch.Tensor, + m_splits: torch.Tensor, + grad_output: torch.Tensor, + *, + recipe: Optional[Recipe], + mode: str, +) -> tuple[torch.Tensor, Optional[torch.Tensor], list[Optional[torch.Tensor]]]: + """Run one microbatch and retain outputs and gradients for path parity.""" + output = _run_microbatches( + layer, + x, + m_splits, + grad_output, + recipe=recipe, + mode=mode, + num_microbatches=1, + ) + input_grad = None if x.grad is None else x.grad.detach().clone() + parameter_grads = [ + None if param.grad is None else param.grad.detach().clone() for param in layer.parameters() + ] + return output.detach().clone(), input_grad, parameter_grads + + +def _validate_path_parity( + *, + multistream_layer: GroupedLinear, + grouped_layer: GroupedLinear, + multistream_x: torch.Tensor, + grouped_x: torch.Tensor, + multistream_splits: torch.Tensor, + grouped_splits: torch.Tensor, + grad_output: torch.Tensor, + recipe: Optional[Recipe], + mode: str, +) -> None: + """Require the two execution paths to produce compatible results.""" + reference = _run_correctness_step( + multistream_layer, + multistream_x, + multistream_splits, + grad_output, + recipe=recipe, + mode=mode, + ) + actual = _run_correctness_step( + grouped_layer, + grouped_x, + grouped_splits, + grad_output, + recipe=recipe, + mode=mode, + ) + + tolerances = {"rtol": 1e-2, "atol": 1e-2} + + torch.testing.assert_close(actual[0], reference[0], **tolerances) + if mode == "fwd": + return + + assert actual[1] is not None and reference[1] is not None + torch.testing.assert_close(actual[1], reference[1], **tolerances) + assert len(actual[2]) == len(reference[2]) + for actual_grad, reference_grad in zip(actual[2], reference[2]): + assert actual_grad is not None and reference_grad is not None + torch.testing.assert_close(actual_grad, reference_grad, **tolerances) + + +def _benchmark_path( + *, + layer: GroupedLinear, + x: torch.Tensor, + m_splits: torch.Tensor, + grad_output: torch.Tensor, + recipe: Optional[Recipe], + mode: str, + num_microbatches: int, + warmup_steps: int, + min_run_time: float, + profile: bool, + label: str, +) -> float: + """Benchmark one execution path and return milliseconds per microbatch.""" + _run_microbatches( + layer, + x, + m_splits, + grad_output, + recipe=recipe, + mode=mode, + num_microbatches=warmup_steps, + ) + torch.cuda.synchronize() + + if profile: + torch.cuda.nvtx.range_push(label) + + timing = benchmark.Timer( + stmt=( + "_run_microbatches(layer, x, m_splits, grad_output, recipe=recipe, " + "mode=mode, num_microbatches=num_microbatches)" + ), + globals={ + "_run_microbatches": _run_microbatches, + "layer": layer, + "x": x, + "m_splits": m_splits, + "grad_output": grad_output, + "recipe": recipe, + "mode": mode, + "num_microbatches": num_microbatches, + }, + num_threads=1, + ).blocked_autorange(min_run_time=min_run_time) + + if profile: + torch.cuda.nvtx.range_pop() + + return timing.median * 1000 / num_microbatches + + +def _gemm_tflops( + *, + total_rows: int, + projection: Projection, + mode: str, + time_ms: float, +) -> float: + """Report GEMM FLOP/s while timing also includes quantization and Python overhead.""" + flops = 2 * total_rows * projection.in_features * projection.out_features + if mode == "fwd_bwd": + flops *= 3 + return flops / time_ms / 1e9 + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--precision", + choices=("all", "bf16", "mxfp8"), + default="all", + ) + parser.add_argument( + "--projection", + choices=("all", "fc1", "fc2"), + default="all", + ) + parser.add_argument( + "--mode", + choices=("fwd", "fwd_bwd"), + default="fwd_bwd", + ) + parser.add_argument("--num-microbatches", type=int, default=16) + parser.add_argument("--warmup-steps", type=int, default=500) + parser.add_argument("--min-run-time", type=float, default=30.0) + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument( + "--m-splits", + type=str, + default=None, + help="Optional comma-separated per-expert rows; defaults to Qwen3.5-397B EP32.", + ) + parser.add_argument( + "--skip-correctness", + action="store_true", + help="Skip output, input-gradient, and parameter-gradient parity checks.", + ) + parser.add_argument("--profile", action="store_true", help="Add per-path NVTX ranges.") + parser.add_argument("--output-csv", type=str, default=None) + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for this benchmark.") + if args.num_microbatches < 1 or args.warmup_steps < 1: + raise ValueError("num_microbatches and warmup_steps must both be positive.") + + if args.m_splits is None: + num_local_experts = QWEN_NUM_EXPERTS // QWEN_EXPERT_PARALLEL_SIZE + rows_per_expert = QWEN_SEQUENCE_LENGTH * QWEN_TOP_K // num_local_experts + split_sizes = [rows_per_expert] * num_local_experts + else: + split_sizes = [int(value) for value in args.m_splits.split(",") if value] + if not split_sizes or any(value < 0 for value in split_sizes): + raise ValueError("m_splits must contain non-negative integers.") + num_local_experts = len(split_sizes) + + total_rows = sum(split_sizes) + if total_rows == 0: + raise ValueError("At least one routed token row is required for this benchmark.") + + precision_names = ("bf16", "mxfp8") if args.precision == "all" else (args.precision,) + projection_names = ("fc1", "fc2") if args.projection == "all" else (args.projection,) + + recipes: dict[str, Optional[Recipe]] = { + "bf16": None, + "mxfp8": MXFP8BlockScaling(), + } + mxfp8_available, reason_for_no_mxfp8 = FP8GlobalStateManager.is_mxfp8_available() + + print("Qwen3.5-397B-A17B GroupedLinear benchmark") + print(f" local experts: {num_local_experts}") + print(f" total routed rows: {total_rows}") + print(f" m_splits: {split_sizes}") + print(f" mode: {args.mode}") + print(" primary parameter dtype: BF16") + print(f" microbatches per timing invocation: {args.num_microbatches}") + print() + + rows = [] + for precision_name in precision_names: + recipe = recipes[precision_name] + if precision_name == "mxfp8" and not mxfp8_available: + print(f"Skipping MXFP8: {reason_for_no_mxfp8}") + continue + if not is_module_grouped_tensor_path_supported(recipe, torch.bfloat16): + print( + f"Skipping {precision_name}: the cuBLASLt grouped-tensor path is unsupported " + "on this GPU or cuBLASLt version." + ) + continue + + for projection_name in projection_names: + projection = PROJECTIONS[projection_name] + torch.manual_seed(args.seed) + torch.cuda.manual_seed(args.seed) + + layers = { + path.name: _build_layer( + projection=projection, + num_local_experts=num_local_experts, + use_grouped_tensor=path.use_grouped_tensor, + ) + for path in EXECUTION_PATHS + } + layers["cublas_grouped_gemm"].load_state_dict(layers["multistream"].state_dict()) + + base_x = torch.randn( + total_rows, + projection.in_features, + dtype=torch.bfloat16, + device="cuda", + ) + inputs = { + path.name: base_x.detach().clone().requires_grad_(args.mode == "fwd_bwd") + for path in EXECUTION_PATHS + } + grad_output = torch.randn( + total_rows, + projection.out_features, + dtype=torch.bfloat16, + device="cuda", + ) + splits = { + path.name: _make_m_splits( + use_grouped_tensor=path.use_grouped_tensor, + split_sizes=split_sizes, + ) + for path in EXECUTION_PATHS + } + + if not args.skip_correctness: + _validate_path_parity( + multistream_layer=layers["multistream"], + grouped_layer=layers["cublas_grouped_gemm"], + multistream_x=inputs["multistream"], + grouped_x=inputs["cublas_grouped_gemm"], + multistream_splits=splits["multistream"], + grouped_splits=splits["cublas_grouped_gemm"], + grad_output=grad_output, + recipe=recipe, + mode=args.mode, + ) + + timings = {} + for path in EXECUTION_PATHS: + label = f"{precision_name}_{projection_name}_{args.mode}_{path.name}" + timing_ms = _benchmark_path( + layer=layers[path.name], + x=inputs[path.name], + m_splits=splits[path.name], + grad_output=grad_output, + recipe=recipe, + mode=args.mode, + num_microbatches=args.num_microbatches, + warmup_steps=args.warmup_steps, + min_run_time=args.min_run_time, + profile=args.profile, + label=label, + ) + timings[path.name] = timing_ms + + speedup = timings["multistream"] / timings["cublas_grouped_gemm"] + for path in EXECUTION_PATHS: + timing_ms = timings[path.name] + rows.append( + { + "precision": precision_name, + "projection": projection_name, + "mode": args.mode, + "execution_path": path.name, + "num_local_experts": num_local_experts, + "total_rows": total_rows, + "in_features": projection.in_features, + "out_features": projection.out_features, + "time_ms": timing_ms, + "gemm_tflops": _gemm_tflops( + total_rows=total_rows, + projection=projection, + mode=args.mode, + time_ms=timing_ms, + ), + "speedup_vs_multistream": (1.0 if path.name == "multistream" else speedup), + } + ) + + print( + f"{precision_name:6s} {projection_name} {args.mode}: " + f"multistream={timings['multistream']:.3f} ms, " + f"cuBLAS grouped={timings['cublas_grouped_gemm']:.3f} ms, " + f"speedup={speedup:.3f}x" + ) + + results = pd.DataFrame(rows) + print() + print(results.to_string(index=False)) + if args.output_csv is not None: + results.to_csv(args.output_csv, index=False) + print(f"\nWrote {args.output_csv}") + + +if __name__ == "__main__": + main() diff --git a/build_tools/VERSION.txt b/build_tools/VERSION.txt index 830a65a39c..e3daa3c6a7 100644 --- a/build_tools/VERSION.txt +++ b/build_tools/VERSION.txt @@ -1 +1 @@ -2.18.0.dev0 +2.19.0.dev0 diff --git a/build_tools/build_ext.py b/build_tools/build_ext.py index 1079accb87..334cfcdf0a 100644 --- a/build_tools/build_ext.py +++ b/build_tools/build_ext.py @@ -6,11 +6,12 @@ """Installation script.""" +import copy import os import subprocess import sys import sysconfig -import copy +import tempfile import time from pathlib import Path @@ -131,23 +132,37 @@ def run(self) -> None: for ext in self.extensions: package_path = Path(self.get_ext_fullpath(ext.name)) install_dir = package_path.resolve().parent - if isinstance(ext, CMakeExtension): - print(f"Building CMake extension {ext.name}") - # Set up incremental builds for CMake extensions - build_dir = os.getenv("NVTE_CMAKE_BUILD_DIR") - if build_dir: - build_dir = Path(build_dir).resolve() - else: - root_dir = Path(__file__).resolve().parent.parent - build_dir = root_dir / "build" / "cmake" - - # Ensure the directory exists + if not isinstance(ext, CMakeExtension): + continue + + print(f"Building CMake extension {ext.name}") + configured_build_dir = os.getenv("NVTE_CMAKE_BUILD_DIR") + if not configured_build_dir and self.inplace: + root_dir = Path(__file__).resolve().parent.parent + configured_build_dir = root_dir / "build" / "cmake" + + if configured_build_dir: + # A persistent build directory enables incremental builds. + build_dir = Path(configured_build_dir).resolve() build_dir.mkdir(parents=True, exist_ok=True) - ext._build_cmake( build_dir=build_dir, install_dir=install_dir, ) + continue + + # Isolate CMake state between concurrent and successive builds. + build_temp = Path(self.build_temp) + build_temp.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory( + prefix=f"cmake-build-{ext.name}-", + dir=build_temp, + ) as build_dir: + print(f"Building CMake extension {ext.name} in temporary directory {build_dir}") + ext._build_cmake( + build_dir=Path(build_dir), + install_dir=install_dir, + ) # Build non-CMake extensions as usual all_extensions = self.extensions diff --git a/docs/envvars.rst b/docs/envvars.rst index b3765a06bd..97eaed5ddc 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -119,6 +119,25 @@ Runtime Environment Variables These environment variables control the behavior of Transformer Engine during execution. +General +^^^^^^^ + +.. envvar:: NVTE_TENSOR_HANDLE_POOL_SIZE_MB + + :Type: ``int`` (positive integer) + :Default: ``20`` + :Description: Size in MiB of the internal ``NVTETensor`` handle pool. Increase this + value if an application legitimately creates more tensor handles than + the default pool can hold. + +.. envvar:: NVTE_GROUPED_TENSOR_HANDLE_POOL_SIZE_MB + + :Type: ``int`` (positive integer) + :Default: ``20`` + :Description: Size in MiB of the internal ``NVTEGroupedTensor`` handle pool. Increase + this value if an application legitimately creates more grouped tensor + handles than the default pool can hold. + Attention Backend Selection ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -360,6 +379,16 @@ Torch Compilation and Fusion LayerNorm/RMSNorm SM Margins ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +.. envvar:: NVTE_CUDNN_MXFP8_NORM_OUTPUT_IN_INPUT_DTYPE + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: With cuDNN 9.25.0 or later, use the normalization input datatype for the virtual + LayerNorm/RMSNorm output consumed by cuDNN MXFP8 block-scale quantization. This + enables cuDNN's fused MXFP8 normalization engine, which requires matching FP16 or + BF16 input and normalization-output datatypes. When set to ``0``, or with an + earlier cuDNN version, the virtual normalization output uses FP32. + .. envvar:: NVTE_FWD_LAYERNORM_SM_MARGIN :Type: ``int`` diff --git a/docs/examples/jax/attention.out b/docs/examples/jax/attention.out new file mode 100644 index 0000000000..a69cdef4e6 --- /dev/null +++ b/docs/examples/jax/attention.out @@ -0,0 +1,13 @@ +# SINGLE_GPU_OUTPUT_START +Native JAX bf16 GQA + SWA: +Mean time: 5.109810829162598 ms + +TE DotProductAttention GQA + SWA: +Mean time: 0.09856224060058594 ms +# SINGLE_GPU_OUTPUT_END + +# MLA_OUTPUT_START +TE DeepSeek-style MLA head dimensions: q/k head dim=128, v head dim=64 +Output shape=(2, 4096, 128, 64), dtype=bfloat16 +Grad shapes=[(2, 4096, 128, 128), (2, 4096, 8, 128), (2, 4096, 8, 64)] +# MLA_OUTPUT_END diff --git a/docs/examples/jax/attention.py b/docs/examples/jax/attention.py new file mode 100644 index 0000000000..3563899332 --- /dev/null +++ b/docs/examples/jax/attention.py @@ -0,0 +1,292 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""JAX: BSHD attention with TransformerEngine. + +Companion source for ``attention.rst``. Code blocks between +``# ATTENTION_*_START`` / ``# ATTENTION_*_END`` markers are pulled into the RST +via ``literalinclude``. + +Run as a script to exercise the example end-to-end: + + python docs/examples/jax/attention.py +""" + +# ATTENTION_IMPORTS_START +from typing import Optional, Tuple + +import jax +import jax.numpy as jnp +import numpy as np +from flax import linen as nn + +import quickstart_jax_utils as utils + +from transformer_engine.jax.attention import SequenceDescriptor +from transformer_engine.jax.flax import DotProductAttention + +# ATTENTION_IMPORTS_END + + +# ATTENTION_INPUTS_START +batch, seq, num_query_heads, num_kv_heads, head_dim = 2, 4096, 128, 8, 128 +window_size = (128, 0) +dtype = jnp.bfloat16 +timing_iters = 20 +warmup_iters = 10 + + +def create_qkv_inputs( + *, + seed: int, + kv_heads: int = num_kv_heads, + qk_head_dim: int = head_dim, + v_head_dim: int = head_dim, +): + """Create separate BSHD query, key, value tensors and an output gradient.""" + + q_key, k_key, v_key, dout_key = jax.random.split(jax.random.PRNGKey(seed), 4) + q = jax.random.normal(q_key, (batch, seq, num_query_heads, qk_head_dim)).astype(dtype) + k = jax.random.normal(k_key, (batch, seq, kv_heads, qk_head_dim)).astype(dtype) + v = jax.random.normal(v_key, (batch, seq, kv_heads, v_head_dim)).astype(dtype) + dout = jax.random.normal(dout_key, (batch, seq, num_query_heads, v_head_dim)).astype(dtype) + return q, k, v, dout + + +def create_full_sequence_descriptor(): + """Describe a BSHD batch with no padding.""" + + seqlens = jnp.full((batch,), seq, dtype=jnp.int32) + return SequenceDescriptor.from_seqlens(seqlens) + + +q, k, v, dout = create_qkv_inputs(seed=2026) +qkv = (q, k, v) +sequence_descriptor = create_full_sequence_descriptor() +# ATTENTION_INPUTS_END + + +# ATTENTION_BASELINE_MODEL_START +def _repeat_kv_for_gqa(x, query_heads): + """Repeat each KV head across its group of query heads.""" + + repeats = query_heads // x.shape[2] + return jnp.repeat(x, repeats, axis=2) + + +def _make_causal_swa_mask(q_len, kv_len, window: Optional[Tuple[int, int]]): + """Create a boolean causal mask, optionally restricted to an SWA window.""" + + q_pos = jnp.arange(q_len)[:, None] + kv_pos = jnp.arange(kv_len)[None, :] + + if window is None: + return kv_pos <= q_pos + + left, right = window + allowed = kv_pos <= q_pos + right + if left >= 0: + allowed = allowed & (kv_pos >= q_pos - left) + return allowed + + +class FlaxNativeGQAAttention(nn.Module): + """Plain JAX/Flax GQA used as the bf16 baseline.""" + + window_size: Optional[Tuple[int, int]] = None + + @nn.compact + def __call__(self, qkv_tensors): + query, key, value = qkv_tensors + key = _repeat_kv_for_gqa(key, query.shape[2]) + value = _repeat_kv_for_gqa(value, query.shape[2]) + + scale = query.shape[-1] ** -0.5 + scores = jnp.einsum( + "bqhd,bkhd->bhqk", + query.astype(jnp.float32), + key.astype(jnp.float32), + ) + scores *= scale + + mask = _make_causal_swa_mask(query.shape[1], key.shape[1], self.window_size) + scores = jnp.where(mask[None, None, :, :], scores, jnp.finfo(jnp.float32).min) + probs = jax.nn.softmax(scores, axis=-1) + out = jnp.einsum("bhqk,bkhd->bqhd", probs, value.astype(jnp.float32)) + return out.astype(query.dtype) + + +baseline = FlaxNativeGQAAttention(window_size=window_size) +baseline_vars = baseline.init(jax.random.PRNGKey(2026), qkv) +# ATTENTION_BASELINE_MODEL_END + + +# ATTENTION_TE_MODEL_START +class TEDotProductAttention(nn.Module): + """Thin Flax wrapper around TE's DotProductAttention.""" + + num_query_heads: int + num_kv_heads: int + qk_head_dim: int = head_dim + attn_mask_type: str = "causal" + qkv_layout: str = "bshd_bshd_bshd" + window_size: Optional[Tuple[int, int]] = None + + @nn.compact + def __call__( + self, + qkv_tensors, + sequence_descriptor: Optional[SequenceDescriptor] = None, + *, + deterministic: bool = False, + ): + query, key, value = qkv_tensors + return DotProductAttention( + head_dim=self.qk_head_dim, + num_attention_heads=self.num_query_heads, + num_gqa_groups=self.num_kv_heads, + attn_mask_type=self.attn_mask_type, + qkv_layout=self.qkv_layout, + attention_dropout=0.0, + transpose_batch_sequence=False, + window_size=self.window_size, + )( + query, + key, + value, + sequence_descriptor=sequence_descriptor, + deterministic=deterministic, + ) + + +te_model = TEDotProductAttention( + num_query_heads=num_query_heads, + num_kv_heads=num_kv_heads, + window_size=window_size, +) +te_vars = te_model.init( + jax.random.PRNGKey(2026), + qkv, + sequence_descriptor=sequence_descriptor, + deterministic=False, +) +# ATTENTION_TE_MODEL_END + + +def run_forward_backward(model, variables, input_qkv, output_grad, seq_desc=None): + """Run one compiled forward+backward pass through an attention module.""" + + def loss_fn(qkv_arg): + if seq_desc is None: + out = model.apply(variables, qkv_arg) + else: + out = model.apply( + variables, + qkv_arg, + sequence_descriptor=seq_desc, + deterministic=False, + ) + return jnp.vdot(out.astype(jnp.float32), output_grad.astype(jnp.float32)) + + return jax.jit(jax.value_and_grad(loss_fn))(input_qkv) + + +def compare_te_to_baseline(input_qkv=qkv, output_grad=dout, seq_desc=sequence_descriptor): + """Compare the TE example to the native baseline.""" + + loss_ref, grads_ref = run_forward_backward(baseline, baseline_vars, input_qkv, output_grad) + loss_te, grads_te = run_forward_backward(te_model, te_vars, input_qkv, output_grad, seq_desc) + out_ref = baseline.apply(baseline_vars, input_qkv) + out_te = te_model.apply(te_vars, input_qkv, sequence_descriptor=seq_desc, deterministic=False) + + jax.block_until_ready((loss_ref, grads_ref, loss_te, grads_te, out_ref, out_te)) + np.testing.assert_allclose(out_te, out_ref, rtol=5e-2, atol=5e-2) + for got, expected in zip(grads_te, grads_ref): + np.testing.assert_allclose(got, expected, rtol=8e-2, atol=8e-2) + + +# ATTENTION_SINGLE_GPU_BENCH_START +def run_single_gpu_bench(): + forward_kwargs = { + "sequence_descriptor": sequence_descriptor, + "deterministic": False, + } + + print("Native JAX bf16 GQA + SWA:") + utils.speedometer( + model_apply_fn=baseline.apply, + variables=baseline_vars, + input=qkv, + output_grad=dout, + timing_iters=timing_iters, + warmup_iters=warmup_iters, + ) + + print("\nTE DotProductAttention GQA + SWA:") + utils.speedometer( + model_apply_fn=te_model.apply, + variables=te_vars, + input=qkv, + output_grad=dout, + forward_kwargs=forward_kwargs, + timing_iters=timing_iters, + warmup_iters=warmup_iters, + ) + + +# ATTENTION_SINGLE_GPU_BENCH_END + + +# ATTENTION_MLA_START +mla_head_dim_qk, mla_head_dim_v = 128, 64 +mla_q, mla_k, mla_v, mla_dout = create_qkv_inputs( + seed=2027, + kv_heads=num_kv_heads, + qk_head_dim=mla_head_dim_qk, + v_head_dim=mla_head_dim_v, +) +mla_qkv = (mla_q, mla_k, mla_v) + +mla_model = TEDotProductAttention( + num_query_heads=num_query_heads, + num_kv_heads=num_kv_heads, + qk_head_dim=mla_head_dim_qk, + window_size=None, +) +mla_vars = mla_model.init( + jax.random.PRNGKey(4), + mla_qkv, + sequence_descriptor=sequence_descriptor, + deterministic=False, +) + + +def run_mla_variant(): + out = mla_model.apply( + mla_vars, + mla_qkv, + sequence_descriptor=sequence_descriptor, + deterministic=False, + ) + loss, grads = run_forward_backward(mla_model, mla_vars, mla_qkv, mla_dout, sequence_descriptor) + jax.block_until_ready((out, loss, grads)) + print( + "TE DeepSeek-style MLA head dimensions: " + f"q/k head dim={mla_head_dim_qk}, v head dim={mla_head_dim_v}" + ) + print(f"Output shape={tuple(out.shape)}, dtype={out.dtype}") + print(f"Grad shapes={[tuple(grad.shape) for grad in grads]}") + + +# ATTENTION_MLA_END + + +if __name__ == "__main__": + print("# SINGLE_GPU_OUTPUT_START") + run_single_gpu_bench() + print("# SINGLE_GPU_OUTPUT_END") + + print("\n# MLA_OUTPUT_START") + run_mla_variant() + print("# MLA_OUTPUT_END") diff --git a/docs/examples/jax/attention.rst b/docs/examples/jax/attention.rst index c9f84da634..f82a6ffe62 100644 --- a/docs/examples/jax/attention.rst +++ b/docs/examples/jax/attention.rst @@ -6,6 +6,44 @@ JAX: Attention with TransformerEngine ===================================== -**TODO — Coming soon.** +Transformer Engine's JAX attention APIs support self-attention and +cross-attention with MHA, GQA, and MQA. Inputs can use standard BSHD or packed +THD layouts, with Q/K/V supplied separately or in packed forms. + +Common options include causal and padding masks, bias, dropout, sliding-window +attention (SWA), attention sinks, and experimental ``score_mod`` callbacks for +FlexAttention-style customization. The API also supports different Q/K and V +head dimensions, as used by the attention operation after the projection +stages in DeepSeek-style MLA. + +For long contexts, selected BSHD and THD configurations support context +parallelism with Ring or AllGather collectives. Exact fused-kernel availability +depends on the input shape, dtype, GPU architecture, and feature combination; +see the `JAX DotProductAttention API reference +<../../api/jax.html#transformer_engine.jax.flax.DotProductAttention>`_ for the +full interface. Choose the tutorial that matches how the sequence dimension is +distributed in your model. `← Back to the JAX integration overview <../te_jax_integration.html>`_ + +Pick a tutorial +--------------- + +.. list-table:: + :header-rows: 1 + :widths: 30, 70 + + * - Tutorial + - Covers + * - `Single-GPU Attention `_ + - BSHD GQA + SWA; performance against a native JAX baseline; + DeepSeek-style MLA head dimensions + * - `Context-Parallel Attention `_ + - Packed THD GQA + SWA on four GPUs; Ring and AllGather CP with striped + load balancing; performance against single-GPU fused attention + +.. toctree:: + :hidden: + + attention_single_gpu + attention_context_parallel diff --git a/docs/examples/jax/attention_context_parallel.out b/docs/examples/jax/attention_context_parallel.out new file mode 100644 index 0000000000..100044be91 --- /dev/null +++ b/docs/examples/jax/attention_context_parallel.out @@ -0,0 +1,16 @@ +# SINGLE_GPU_OUTPUT_START +Single-GPU THD GQA + SWA: +Mean time: 126.68747901916504 ms +# SINGLE_GPU_OUTPUT_END + +# RING_OUTPUT_START +THD CP Ring stripe_size=1: +Mean time: 57.16729164123535 ms +Speedup vs single GPU: 2.22x +# RING_OUTPUT_END + +# AG_OUTPUT_START +THD CP AllGather stripe_size=512: +Mean time: 53.79219055175781 ms +Speedup vs single GPU: 2.36x +# AG_OUTPUT_END diff --git a/docs/examples/jax/attention_context_parallel.py b/docs/examples/jax/attention_context_parallel.py new file mode 100644 index 0000000000..1557a30b7c --- /dev/null +++ b/docs/examples/jax/attention_context_parallel.py @@ -0,0 +1,417 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""JAX: context-parallel THD attention with Transformer Engine. + +Companion source for ``attention_context_parallel.rst``. Code blocks between +``# ATTENTION_CP_*_START`` / ``# ATTENTION_CP_*_END`` markers are pulled into +the RST via ``literalinclude``. + +Run as a script to exercise the example end-to-end: + + python docs/examples/jax/attention_context_parallel.py +""" + +# ATTENTION_CP_IMPORTS_START +import os +import time +from typing import Tuple + +# Ring + SWA uses the non-scan Ring implementation. Set this before JAX compiles +# the first fused attention call so the example follows the distributed tests. +os.environ.setdefault("NVTE_FUSED_RING_ATTENTION_USE_SCAN", "0") + +import jax +import jax.numpy as jnp +import numpy as np +from jax.sharding import Mesh, NamedSharding, PartitionSpec as P + +import transformer_engine.jax as te +from transformer_engine.jax.attention import ( + AttnBiasType, + AttnMaskType, + AttnSoftmaxType, + CPStrategy, + QKVLayout, + ReorderStrategy, + SequenceDescriptor, + fused_attn, + inverse_reorder_causal_load_balancing, + is_fused_attn_kernel_available, + reorder_causal_load_balancing, +) +from transformer_engine.jax.sharding import MeshResource + +# ATTENTION_CP_IMPORTS_END + + +# ATTENTION_CP_INPUTS_START +cp_size = 4 +batch, seq, num_query_heads, num_kv_heads, head_dim = 2, 65536, 128, 8, 128 +runtime_segments_per_seq = 4 +max_segments_per_seq = runtime_segments_per_seq +window_size = (8192, 0) +dtype = jnp.bfloat16 +timing_iters = 5 +warmup_iters = 2 +ring_stripe_size = 1 +ag_stripe_size = 512 + + +def create_qkv_inputs(seed: int = 2026): + """Create separate THD GQA tensors and an output gradient.""" + + q_key, k_key, v_key, dout_key = jax.random.split(jax.random.PRNGKey(seed), 4) + q_shape = (batch, seq, num_query_heads, head_dim) + kv_shape = (batch, seq, num_kv_heads, head_dim) + q = jax.random.normal(q_key, q_shape).astype(dtype) + k = jax.random.normal(k_key, kv_shape).astype(dtype) + v = jax.random.normal(v_key, kv_shape).astype(dtype) + dout = jax.random.normal(dout_key, q_shape).astype(dtype) + return q, k, v, dout + + +def create_packed_segment_ids_and_pos(): + """Pack padded causal segments into each THD batch row.""" + + segment_slot_len = seq // runtime_segments_per_seq + valid_segment_len = 3 * segment_slot_len // 4 + segment_ids_per_row = [] + segment_pos_per_row = [] + + for segment_id in range(1, runtime_segments_per_seq + 1): + valid_ids = jnp.full((valid_segment_len,), segment_id, dtype=jnp.int32) + padded_ids = jnp.zeros((segment_slot_len - valid_segment_len,), dtype=jnp.int32) + segment_ids_per_row.append(jnp.concatenate([valid_ids, padded_ids])) + segment_pos_per_row.append(jnp.arange(segment_slot_len, dtype=jnp.int32)) + + segment_ids = jnp.concatenate(segment_ids_per_row) + segment_pos = jnp.concatenate(segment_pos_per_row) + segment_ids = jnp.tile(segment_ids[None, :], (batch, 1)) + segment_pos = jnp.tile(segment_pos[None, :], (batch, 1)) + return segment_ids, segment_pos + + +def create_sequence_descriptor(segment_ids_arg, segment_pos_arg): + """Create the THD sequence descriptor from segment IDs and positions.""" + + return SequenceDescriptor.from_segment_ids_and_pos(segment_ids_arg, segment_pos_arg) + + +q, k, v, dout = create_qkv_inputs() +segment_ids, segment_pos = create_packed_segment_ids_and_pos() +sequence_descriptor = create_sequence_descriptor(segment_ids, segment_pos) +# ATTENTION_CP_INPUTS_END + + +# ATTENTION_CP_MESH_START +def build_cp_mesh(): + """Use one JAX mesh axis for context parallelism over sequence.""" + + devices = np.asarray(jax.devices()[:cp_size]) + mesh = Mesh(devices, axis_names=("cp",)) + # Also set the corresponding MeshResource fields when other parallelisms + # use additional mesh axis names in the surrounding model. + mesh_resource = MeshResource(cp_resource="cp") + return mesh, mesh_resource + + +# ATTENTION_CP_MESH_END + + +# ATTENTION_CP_FUSED_ATTENTION_START +def fused_thd_attention( + qkv_tensors, + seq_desc, + *, + context_parallel_axis: str = "", + context_parallel_strategy: CPStrategy = CPStrategy.DEFAULT, + context_parallel_causal_load_balanced: bool = False, + stripe_size: int | None = None, +): + """Call TE fused attention on separate THD Q, K, V tensors.""" + + return fused_attn( + qkv_tensors, + None, + seq_desc, + None, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.PADDING_CAUSAL_MASK, + qkv_layout=QKVLayout.THD_THD_THD, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + scaling_factor=head_dim**-0.5, + dropout_probability=0.0, + is_training=True, + max_segments_per_seq=max_segments_per_seq, + window_size=window_size, + context_parallel_strategy=context_parallel_strategy, + context_parallel_causal_load_balanced=context_parallel_causal_load_balanced, + context_parallel_axis=context_parallel_axis, + stripe_size=stripe_size, + ) + + +def apply_context_parallel_attention( + _variables, + qkv_tensors, + *, + seq_desc, + context_parallel_strategy: CPStrategy, + stripe_size: int, + rngs=None, +): + del rngs + return fused_thd_attention( + qkv_tensors, + seq_desc, + context_parallel_axis="cp", + context_parallel_strategy=context_parallel_strategy, + context_parallel_causal_load_balanced=True, + stripe_size=stripe_size, + ) + + +# ATTENTION_CP_FUSED_ATTENTION_END + + +# ATTENTION_CP_REORDER_START +def reorder_for_context_parallel(x, stripe_size: int): + return reorder_causal_load_balancing( + x, + strategy=ReorderStrategy.Striped, + cp_size=cp_size, + seq_dim=1, + stripe_size=stripe_size, + ) + + +def inverse_reorder_from_context_parallel(x, stripe_size: int): + return inverse_reorder_causal_load_balancing( + x, + strategy=ReorderStrategy.Striped, + cp_size=cp_size, + seq_dim=1, + stripe_size=stripe_size, + ) + + +def create_reordered_sequence_descriptor(stripe_size: int): + reordered_ids = reorder_for_context_parallel(segment_ids, stripe_size) + reordered_pos = reorder_for_context_parallel(segment_pos, stripe_size) + return create_sequence_descriptor(reordered_ids, reordered_pos) + + +# ATTENTION_CP_REORDER_END + + +# ATTENTION_CP_SHARD_START +def shard_sequence_descriptor(mesh, seq_desc): + def put_leaf(x): + if x.ndim == 1: + sharding = NamedSharding(mesh, P(None)) + else: + sharding = NamedSharding(mesh, P(None, "cp")) + return jax.device_put(x, sharding) + + return jax.tree.map(put_leaf, seq_desc) + + +def shard_for_context_parallel(mesh, stripe_size: int): + qkv_sharding = NamedSharding(mesh, P(None, "cp", None, None)) + dout_sharding = NamedSharding(mesh, P(None, "cp", None, None)) + reordered_seq_desc = create_reordered_sequence_descriptor(stripe_size) + + return { + "qkv": tuple( + jax.device_put(reorder_for_context_parallel(x, stripe_size), qkv_sharding) + for x in (q, k, v) + ), + "dout": jax.device_put(reorder_for_context_parallel(dout, stripe_size), dout_sharding), + "sequence_descriptor": shard_sequence_descriptor(mesh, reordered_seq_desc), + } + + +# ATTENTION_CP_SHARD_END + + +def _strategy_name(strategy: CPStrategy): + return "Ring" if strategy == CPStrategy.RING else "AllGather" + + +def context_parallel_supported() -> Tuple[bool, str]: + if len(jax.devices()) < cp_size: + return False, f"needs {cp_size} GPUs" + + has_kernel = is_fused_attn_kernel_available( + True, + dtype, + dtype, + QKVLayout.THD_THD_THD, + AttnBiasType.NO_BIAS, + AttnMaskType.PADDING_CAUSAL_MASK, + AttnSoftmaxType.VANILLA_SOFTMAX, + 0.0, + num_query_heads, + num_kv_heads, + seq, + seq, + head_dim, + head_dim, + window_size, + ) + if not has_kernel: + return False, "no fused attention kernel for the THD SWA shape" + return True, "" + + +def _single_gpu_grad_fn(): + def loss_fn(qkv_arg, seq_desc_arg, dout_arg): + out = fused_thd_attention(qkv_arg, seq_desc_arg) + return jnp.vdot(out.astype(jnp.float32), dout_arg.astype(jnp.float32)) + + return jax.jit(jax.value_and_grad(loss_fn)) + + +def run_reference_attention(): + """Run single-GPU fused attention for CP output comparisons.""" + + out = fused_thd_attention((q, k, v), sequence_descriptor) + return jax.block_until_ready(out) + + +# ATTENTION_CP_RUN_START +def _context_parallel_jit_fns(strategy: CPStrategy, stripe_size: int, sharded): + qkv_shardings = tuple(x.sharding for x in sharded["qkv"]) + seq_desc_shardings = jax.tree.map(lambda x: x.sharding, sharded["sequence_descriptor"]) + dout_sharding = sharded["dout"].sharding + + def loss_fn(qkv_arg, seq_desc_arg, dout_arg): + out = apply_context_parallel_attention( + {}, + qkv_arg, + seq_desc=seq_desc_arg, + context_parallel_strategy=strategy, + stripe_size=stripe_size, + ) + return jnp.vdot(out.astype(jnp.float32), dout_arg.astype(jnp.float32)) + + def forward_fn(qkv_arg, seq_desc_arg): + out = apply_context_parallel_attention( + {}, + qkv_arg, + seq_desc=seq_desc_arg, + context_parallel_strategy=strategy, + stripe_size=stripe_size, + ) + return inverse_reorder_from_context_parallel(out, stripe_size) + + grad_fn = jax.jit( + jax.value_and_grad(loss_fn), + in_shardings=(qkv_shardings, seq_desc_shardings, dout_sharding), + out_shardings=(None, qkv_shardings), + ) + forward_jit = jax.jit( + forward_fn, + in_shardings=(qkv_shardings, seq_desc_shardings), + ) + return grad_fn, forward_jit + + +def run_context_parallel_case(strategy: CPStrategy, stripe_size: int): + mesh, mesh_resource = build_cp_mesh() + sharded = shard_for_context_parallel(mesh, stripe_size) + grad_fn, forward_jit = _context_parallel_jit_fns(strategy, stripe_size, sharded) + + with jax.set_mesh(mesh), te.autocast(mesh_resource=mesh_resource): + loss, grads = grad_fn( + sharded["qkv"], + sharded["sequence_descriptor"], + sharded["dout"], + ) + out = forward_jit(sharded["qkv"], sharded["sequence_descriptor"]) + + jax.block_until_ready((loss, grads, out)) + return {"loss": loss, "grads": grads, "output": out} + + +def run_single_gpu_bench(): + grad_fn = _single_gpu_grad_fn() + + print("Single-GPU THD GQA + SWA:") + for _ in range(warmup_iters): + result = grad_fn((q, k, v), sequence_descriptor, dout) + jax.block_until_ready(result) + + start = time.time() + for _ in range(timing_iters): + result = grad_fn((q, k, v), sequence_descriptor, dout) + jax.block_until_ready(result) + mean_ms = (time.time() - start) * 1000 / timing_iters + print(f"Mean time: {mean_ms} ms") + return mean_ms + + +def run_context_parallel_bench( + strategy: CPStrategy, + stripe_size: int, + single_gpu_ms: float | None = None, +): + mesh, mesh_resource = build_cp_mesh() + sharded = shard_for_context_parallel(mesh, stripe_size) + grad_fn, _ = _context_parallel_jit_fns(strategy, stripe_size, sharded) + + print(f"THD CP {_strategy_name(strategy)} stripe_size={stripe_size}:") + with jax.set_mesh(mesh), te.autocast(mesh_resource=mesh_resource): + for _ in range(warmup_iters): + result = grad_fn( + sharded["qkv"], + sharded["sequence_descriptor"], + sharded["dout"], + ) + jax.block_until_ready(result) + + start = time.time() + for _ in range(timing_iters): + result = grad_fn( + sharded["qkv"], + sharded["sequence_descriptor"], + sharded["dout"], + ) + jax.block_until_ready(result) + end = time.time() + + mean_ms = (end - start) * 1000 / timing_iters + print(f"Mean time: {mean_ms} ms") + if single_gpu_ms is not None: + print(f"Speedup vs single GPU: {single_gpu_ms / mean_ms:.2f}x") + + +# ATTENTION_CP_RUN_END + + +if __name__ == "__main__": + supported, reason = context_parallel_supported() + if not supported: + print(f"skipped context-parallel example: {reason}") + else: + print("# SINGLE_GPU_OUTPUT_START") + single_gpu_ms = run_single_gpu_bench() + print("# SINGLE_GPU_OUTPUT_END") + + print("# RING_OUTPUT_START") + run_context_parallel_bench( + CPStrategy.RING, + ring_stripe_size, + single_gpu_ms, + ) + print("# RING_OUTPUT_END") + + print("\n# AG_OUTPUT_START") + run_context_parallel_bench( + CPStrategy.ALL_GATHER, + ag_stripe_size, + single_gpu_ms, + ) + print("# AG_OUTPUT_END") diff --git a/docs/examples/jax/attention_context_parallel.rst b/docs/examples/jax/attention_context_parallel.rst new file mode 100644 index 0000000000..5c2e2dd7ec --- /dev/null +++ b/docs/examples/jax/attention_context_parallel.rst @@ -0,0 +1,181 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +JAX: Context-Parallel Attention with TransformerEngine +====================================================== + +Transformer Engine fused attention supports context parallelism (CP) for +selected BSHD and packed THD Q/K/V layouts; see the +`JAX DotProductAttention API reference +<../../api/jax.html#transformer_engine.jax.flax.DotProductAttention>`_ for the +layout definitions and full interface. This tutorial focuses on a +representative packed THD configuration with +`grouped-query attention (GQA) `_, padded +segments, causal +`sliding-window attention (SWA) `_, and both +Ring and AllGather strategies. + +CP shards the sequence dimension over a JAX mesh axis so long-context attention +can split activation memory and attention work across devices while +Transformer Engine (TE) runs the required collectives inside the fused +attention call. + +.. note:: + + CP is most useful when attention does not fit on one GPU, or when long + sequences and sufficiently wide attention windows provide enough + computation to amortize communication. For instance, applications that use + GQA may be good candidates for CP because GQA's lower K/V head count reduces + communication across devices. Conversely, workloads with narrow SWA windows + may be better suited to single-GPU fused attention: CP still communicates + K/V across devices while each query attends to relatively few tokens. + +**Prerequisite:** this example requires four GPUs. + +`← Back to the Attention overview `_ + +1. Packed THD inputs +-------------------- + +In the separate-QKV THD layout used here, Q/K/V are shaped +``[batch, seq, heads, dim]``, and the sequence dimension can pack several +shorter segments. The ``SequenceDescriptor`` tells TE which tokens belong to +which packed segment and which token slots are padding. CP supports separate +Q/K/V (``THD_THD_THD``) and packed K/V (``THD_T2HD``) layouts, but not fully +packed QKV (``T3HD``); this tutorial uses separate tensors. It uses a batch of +two 64k sequences. Each sequence contains four padded, 16k-capacity segment +slots with 12,288 valid tokens and 4,096 padding tokens per slot. It also uses +GQA with 128 query heads and 8 K/V heads. + +.. literalinclude:: attention_context_parallel.py + :language: python + :start-after: # ATTENTION_CP_IMPORTS_START + :end-before: # ATTENTION_CP_IMPORTS_END + +The tensor inputs and packed-sequence metadata are created as follows. + +.. literalinclude:: attention_context_parallel.py + :language: python + :start-after: # ATTENTION_CP_INPUTS_START + :end-before: # ATTENTION_CP_INPUTS_END + + +2. Context-parallel mesh +------------------------ + +The JAX ``Mesh`` describes the physical devices. ``MeshResource`` tells TE which +mesh axis is used for context parallelism. + +.. literalinclude:: attention_context_parallel.py + :language: python + :start-after: # ATTENTION_CP_MESH_START + :end-before: # ATTENTION_CP_MESH_END + + +3. Fused attention call +----------------------- + +This example calls ``transformer_engine.jax.attention.fused_attn`` directly. The +Flax ``DotProductAttention`` wrapper covers the common path, but the lower-level +function exposes ``stripe_size``. + +.. literalinclude:: attention_context_parallel.py + :language: python + :start-after: # ATTENTION_CP_FUSED_ATTENTION_START + :end-before: # ATTENTION_CP_FUSED_ATTENTION_END + + +4. Striped load balancing and sharding +-------------------------------------- + +For THD causal CP, TE uses striped load balancing. Ring attention requires +``stripe_size=1``. AllGather can use a larger stripe size; this tutorial uses +``stripe_size=512`` for the 64k sequence shape. Ring + SWA uses the non-scan +Ring path, set in the example before the first fused attention call is compiled. + +.. literalinclude:: attention_context_parallel.py + :language: python + :start-after: # ATTENTION_CP_REORDER_START + :end-before: # ATTENTION_CP_REORDER_END + +.. literalinclude:: attention_context_parallel.py + :language: python + :start-after: # ATTENTION_CP_SHARD_START + :end-before: # ATTENTION_CP_SHARD_END + + +5. Ring and AllGather +--------------------- + +The single-GPU baseline and both CP examples use the same packed THD GQA shape, +causal masking, 8192-token SWA window, and dropout-free fused attention. The +only strategy-specific difference between the two CP cases is the strategy and +stripe size. CP collectives depend on the compiler seeing the intended sharding, +so the forward and forward+backward functions are compiled with explicit +``in_shardings``; the forward+backward path also pins the gradient sharding. +The timing loop follows the same forward+backward pattern as ``speedometer`` +while keeping those sharding controls visible. + +.. literalinclude:: attention_context_parallel.py + :language: python + :start-after: # ATTENTION_CP_RUN_START + :end-before: # ATTENTION_CP_RUN_END + +.. raw:: html + +
+ Single-GPU output: +
+ +.. container:: program-output + + .. literalinclude:: attention_context_parallel.out + :language: text + :start-after: # SINGLE_GPU_OUTPUT_START + :end-before: # SINGLE_GPU_OUTPUT_END + +.. raw:: html + +
+ Ring output: +
+ +.. container:: program-output + + .. literalinclude:: attention_context_parallel.out + :language: text + :start-after: # RING_OUTPUT_START + :end-before: # RING_OUTPUT_END + +.. raw:: html + +
+ AllGather output: +
+ +.. container:: program-output + + .. literalinclude:: attention_context_parallel.out + :language: text + :start-after: # AG_OUTPUT_START + :end-before: # AG_OUTPUT_END + +On four GB200s, Ring is roughly **2.22x faster** and AllGather roughly **2.36x +faster** than the equivalent single-GPU fused-attention forward+backward pass. +These results are specific to this workload and system. Applications with long +segments and wide attention windows generally have more attention computation +relative to communication and are stronger CP candidates; workloads with short +windows may see less benefit. Performance also depends on the batch and segment +lengths, head configuration, CP strategy, stripe size, and interconnect. + + + +Next steps +---------- + +* `Single-GPU attention `_: BSHD GQA, SWA, and + DeepSeek-style MLA head dimensions. +* `← Attention overview `_ +* `← Hub <../te_jax_integration.html>`_ diff --git a/docs/examples/jax/attention_single_gpu.rst b/docs/examples/jax/attention_single_gpu.rst new file mode 100644 index 0000000000..f675ebbe05 --- /dev/null +++ b/docs/examples/jax/attention_single_gpu.rst @@ -0,0 +1,149 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +JAX: Single-GPU Attention with TransformerEngine +================================================ + +This document walks through replacing a plain JAX implementation of BSHD +attention with TransformerEngine's fused ``DotProductAttention``. +The example uses +`grouped-query attention (GQA) `_ and +`sliding-window attention (SWA) `_. + +`← Back to the Attention overview `_ + +1. Baseline: native JAX BSHD GQA + SWA +-------------------------------------- + +Start with the imports shared by the native JAX and Transformer Engine +implementations. + +.. literalinclude:: attention.py + :language: python + :start-after: # ATTENTION_IMPORTS_START + :end-before: # ATTENTION_IMPORTS_END + +Next, create reproducible BSHD Q/K/V tensors and the sequence descriptor. +The ``SequenceDescriptor`` supplies TE with sequence lengths and, for packed +inputs, segment boundaries and padding metadata. + +.. literalinclude:: attention.py + :language: python + :start-after: # ATTENTION_INPUTS_START + :end-before: # ATTENTION_INPUTS_END + +The native JAX baseline repeats K/V heads for GQA and applies the causal +sliding-window mask explicitly. + +.. literalinclude:: attention.py + :language: python + :start-after: # ATTENTION_BASELINE_MODEL_START + :end-before: # ATTENTION_BASELINE_MODEL_END + + +2. Transformer Engine ``DotProductAttention`` +---------------------------------------------- + +The Transformer Engine version keeps the same separate BSHD inputs. The important arguments are +``num_gqa_groups`` for GQA, ``attn_mask_type="causal"`` for autoregressive +attention, and ``window_size`` for SWA. + +.. literalinclude:: attention.py + :language: python + :start-after: # ATTENTION_TE_MODEL_START + :end-before: # ATTENTION_TE_MODEL_END + + +3. Single-GPU performance +------------------------- + +``speedometer`` runs a JIT-compiled forward+backward loop with warmup for both +implementations. + +.. literalinclude:: attention.py + :language: python + :start-after: # ATTENTION_SINGLE_GPU_BENCH_START + :end-before: # ATTENTION_SINGLE_GPU_BENCH_END + +.. raw:: html + +
+ Output: +
+ +.. container:: program-output + + .. literalinclude:: attention.out + :language: text + :start-after: # SINGLE_GPU_OUTPUT_START + :end-before: # SINGLE_GPU_OUTPUT_END + +On a single GB200, this run is roughly **52x faster** for the fwd+bwd of this +BSHD GQA + SWA example. This compares TE ``DotProductAttention`` against the +native JAX baseline above, which materializes attention scores with XLA ops; it +is not a comparison against ``jax.nn.dot_product_attention(..., +implementation="cudnn")``. + + +4. DeepSeek-style MLA head dimensions +------------------------------------- + +This example covers the attention-kernel interface used after +`DeepSeek-style MLA projections `_, not the +latent projection layers themselves. At this point, separate Q, K, and V +tensors can use different per-head dimensions for Q/K and V. Keep +``qkv_layout="bshd_bshd_bshd"`` so TE can see the Q/K head dimension and the V +head dimension separately. + +.. literalinclude:: attention.py + :language: python + :start-after: # ATTENTION_MLA_START + :end-before: # ATTENTION_MLA_END + +.. raw:: html + +
+ Output: +
+ +.. container:: program-output + + .. literalinclude:: attention.out + :language: text + :start-after: # MLA_OUTPUT_START + :end-before: # MLA_OUTPUT_END + + +Other attention knobs +--------------------- + +The examples above represent a subset of attention features. Other +``DotProductAttention`` features can be enabled through the same module +arguments as below: + +* Dropout: set ``attention_dropout > 0``, call with ``deterministic=False``, and + pass a Flax ``dropout`` RNG to ``apply``. +* Bias: pass ``bias`` and set ``attn_bias_type`` when the selected fused kernel + supports that bias mode. +* Sink attention: use ``softmax_type="off_by_one"`` or ``"learnable"``. +* Score scaling: set ``scale_factor`` to override the default + ``1 / sqrt(head_dim)`` scaling. +* Determinism: set ``NVTE_ALLOW_NONDETERMINISTIC_ALGO=0`` before launching the + process if deterministic fused kernels are required. +* Score modification (experimental): use ``score_mod`` for a + FlexAttention-style cuDNN frontend callback, with runtime operands in + ``score_mod_tensors`` and optional custom backward logic in + ``score_mod_bprop``. This path requires fused attention and currently cannot + be combined with masks, bias, dropout, SWA, CP, or packed/ragged sequence + metadata. + + +Next steps +---------- + +* `Context-parallel attention `_: packed THD + attention over a context-parallel mesh. +* `← Attention overview `_ +* `← Hub <../te_jax_integration.html>`_ diff --git a/docs/examples/jax/test_attention.py b/docs/examples/jax/test_attention.py new file mode 100644 index 0000000000..3cc08271dc --- /dev/null +++ b/docs/examples/jax/test_attention.py @@ -0,0 +1,191 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Pytest entry points for the JAX attention tutorials.""" + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from transformer_engine.jax.attention import ( + AttnBiasType, + AttnMaskType, + AttnSoftmaxType, + QKVLayout, + is_fused_attn_kernel_available, +) +from transformer_engine_jax import get_device_compute_capability + +# Imports from ``attention`` and ``attention_context_parallel`` are intentionally +# deferred into each test body. The tutorial modules create tensors and initialize +# models at module scope; deferring imports lets pytest apply skip marks before +# unsupported CI nodes allocate those examples. + +requires_hopper_or_newer = pytest.mark.skipif( + get_device_compute_capability(0) < 90, + reason="the 4K native JAX baseline requires more than 48 GiB of device memory", +) + + +def test_bshd_gqa_swa_runs(): + import attention + + out = attention.te_model.apply( + attention.te_vars, + attention.qkv, + sequence_descriptor=attention.sequence_descriptor, + deterministic=False, + ) + + assert out.shape == attention.dout.shape + assert out.dtype == attention.dtype + + +@requires_hopper_or_newer +def test_bshd_gqa_swa_matches_baseline(): + import attention + + attention.compare_te_to_baseline() + + +@requires_hopper_or_newer +def test_single_gpu_benchmark(): + import attention + + attention.run_single_gpu_bench() + + +def test_mla_variant_runs(): + import attention + + out = attention.mla_model.apply( + attention.mla_vars, + attention.mla_qkv, + sequence_descriptor=attention.sequence_descriptor, + deterministic=False, + ) + loss, grads = attention.run_forward_backward( + attention.mla_model, + attention.mla_vars, + attention.mla_qkv, + attention.mla_dout, + attention.sequence_descriptor, + ) + jax.block_until_ready((out, loss, grads)) + + assert out.shape == attention.mla_dout.shape + assert out.dtype == attention.dtype + assert loss.shape == () + assert [grad.shape for grad in grads] == [x.shape for x in attention.mla_qkv] + + +def _context_parallel_supported(): + cp_size = 4 + if len(jax.devices()) < cp_size: + return False, f"needs {cp_size} GPUs" + + has_kernel = is_fused_attn_kernel_available( + True, + jnp.bfloat16, + jnp.bfloat16, + QKVLayout.THD_THD_THD, + AttnBiasType.NO_BIAS, + AttnMaskType.PADDING_CAUSAL_MASK, + AttnSoftmaxType.VANILLA_SOFTMAX, + 0.0, + 128, + 8, + 65536, + 65536, + 128, + 128, + (8192, 0), + ) + if not has_kernel: + return False, "no fused attention kernel for the THD SWA shape" + return True, "" + + +_cp_supported, _cp_reason = _context_parallel_supported() +requires_cp = pytest.mark.skipif( + not _cp_supported, + reason=f"context-parallel attention tutorial skipped: {_cp_reason}", +) + + +def _assert_cp_result(cp_attention, strategy, stripe_size): + result = cp_attention.run_context_parallel_case(strategy, stripe_size) + reference = cp_attention.run_reference_attention() + + assert result["output"].shape == ( + cp_attention.batch, + cp_attention.seq, + cp_attention.num_query_heads, + cp_attention.head_dim, + ) + assert result["output"].dtype == cp_attention.dtype + assert result["loss"].shape == () + assert [grad.shape for grad in result["grads"]] == [ + x.shape for x in cp_attention.create_qkv_inputs()[:3] + ] + + valid_tokens = cp_attention.segment_ids.astype(bool)[..., None, None] + valid_diff = jax.numpy.max( + jax.numpy.where( + valid_tokens, + jax.numpy.abs( + result["output"].astype(jax.numpy.float32) - reference.astype(jax.numpy.float32) + ), + 0.0, + ) + ) + padded_max = jax.numpy.max( + jax.numpy.where( + valid_tokens, + 0.0, + jax.numpy.abs(result["output"].astype(jax.numpy.float32)), + ) + ) + np.testing.assert_allclose(valid_diff, 0, rtol=5e-2, atol=5e-2) + np.testing.assert_allclose(padded_max, 0, rtol=5e-2, atol=5e-2) + + +@requires_cp +def test_multi_gpu_context_parallel_ring_case(): + import attention_context_parallel as cp_attention + + _assert_cp_result( + cp_attention, + cp_attention.CPStrategy.RING, + cp_attention.ring_stripe_size, + ) + + +@requires_cp +def test_multi_gpu_context_parallel_allgather_case(): + import attention_context_parallel as cp_attention + + _assert_cp_result( + cp_attention, + cp_attention.CPStrategy.ALL_GATHER, + cp_attention.ag_stripe_size, + ) + + +@requires_cp +def test_multi_gpu_context_parallel_benchmarks(): + import attention_context_parallel as cp_attention + + single_gpu_ms = cp_attention.run_single_gpu_bench() + cp_attention.run_context_parallel_bench( + cp_attention.CPStrategy.RING, + cp_attention.ring_stripe_size, + single_gpu_ms, + ) + cp_attention.run_context_parallel_bench( + cp_attention.CPStrategy.ALL_GATHER, + cp_attention.ag_stripe_size, + single_gpu_ms, + ) diff --git a/docs/examples/op_fuser/op_fuser.rst b/docs/examples/op_fuser/op_fuser.rst index dd17191e58..a6a500f20e 100644 --- a/docs/examples/op_fuser/op_fuser.rst +++ b/docs/examples/op_fuser/op_fuser.rst @@ -151,6 +151,159 @@ arguments and the extra outputs will be returned. the block has been split into two sections, each with one branching operation. +Extra tensor channels +""""""""""""""""""""" + +Branching operations can also route their extra inputs and outputs within +the same ``Sequential`` via named channels. Extra output tensors with a +specified channel can be consumed by later operations in the same ``Sequential`` +and may optionally be returned to the caller. Extra input tensors with a specified +channel are accessed internally instead of being provided as arguments +to ``Sequential``. + +With a channel, the residual block above can be expressed using one +``Sequential``: + +.. code-block:: python + + import torch + import transformer_engine.pytorch as te + + make_residual = te.ops.MakeExtraOutput() + add_residual = te.ops.AddExtraInput() + make_residual.set_extra_output_channel( + 0, "residual", output_to_caller=False + ) + add_residual.set_extra_input_channel(0, "residual") + + block = te.ops.Sequential( + te.ops.LayerNorm(4096), + make_residual, + te.ops.Linear(4096, 28672), + te.ops.SwiGLU(), + te.ops.Linear(14336, 4096), + add_residual, + ) + + # The residual is routed internally and omitted from the public outputs. + x = torch.randn(16384, 4096, device="cuda") + y = block(x) + +Channels are also useful for mixture-of-experts blocks. The following +example assumes custom ``Dispatch`` and ``Combine`` basic operations. +``Dispatch`` has one public extra input containing router probabilities +and three extra outputs: split sizes, token probabilities, and a +routing map. ``Combine`` consumes the routing map. + +.. code-block:: python + + import transformer_engine.pytorch as te + from my_ops import Dispatch, Combine + + num_experts = 8 + hidden_size = 4096 + ffn_size = 14336 + + dispatch = Dispatch(num_experts) + fc1 = te.ops.GroupedLinear( + num_experts, hidden_size, 2 * ffn_size, bias=False + ) + activation = te.ops.ScaledSwiGLU() + fc2 = te.ops.GroupedLinear( + num_experts, ffn_size, hidden_size, bias=False + ) + combine = Combine(num_experts) + + # Dispatch extra outputs: + # 0: split sizes, 1: token probabilities, 2: routing map + dispatch.set_extra_output_channel( + 0, "m_splits", output_to_caller=False + ) + dispatch.set_extra_output_channel( + 1, "probs", output_to_caller=False + ) + dispatch.set_extra_output_channel( + 2, "routing_map", output_to_caller=False + ) + + fc1.set_extra_input_channel(0, "m_splits") + activation.set_extra_input_channel(0, "probs") + fc2.set_extra_input_channel(0, "m_splits") + combine.set_extra_input_channel(0, "routing_map") + + moe = te.ops.Sequential(dispatch, fc1, activation, fc2, combine) + + # Dispatch's extra input has no channel, so the caller passes router_probs. + # Channels supply all later extra inputs internally. The channel outputs + # are not returned because output_to_caller=False. + y = moe(x, router_probs) + +Channels cannot connect operations in different ``OperationFuser`` +instances. In particular, an ordinary PyTorch module inside a +``Sequential`` splits the fusible operations on either side into +separate fusers. The following channel connection is therefore not +supported: + +.. code-block:: python + + make_residual = te.ops.MakeExtraOutput() + add_residual = te.ops.AddExtraInput() + make_residual.set_extra_output_channel(0, "residual") + add_residual.set_extra_input_channel(0, "residual") + + block = te.ops.Sequential( + make_residual, + torch.nn.Identity(), # Splits the operations into separate fusers. + add_residual, + ) + +Use the public extra output and extra input interfaces, as in the +two-``Sequential`` example above, when the producer and consumer cannot +be placed in the same ``OperationFuser``. + +The following conditions apply to extra tensor channels: + +- Every named extra input must have a matching producer earlier in the + same fuser. Leave an extra input unnamed when the caller should provide it. +- An output channel name has at most one producer, but its output may + fan out to multiple consumers. +- A named output does not require a consumer. It is returned as a public + extra output by default. +- A channel is scoped to one ``OperationFuser``. In a ``Sequential``, + ordinary PyTorch modules split adjacent fusible operations into + separate fusers, and channels cannot cross that boundary. +- The caller passes unnamed extra inputs. Named, channel-connected extra + input slots do not appear in the ``Sequential`` arguments. +- ``set_extra_output_channel`` accepts ``output_to_caller`` (``True`` by + default). Public extra outputs are returned in their original + basic-operation and slot order. Gradients supplied for a returned output + are combined with gradients from its internal channel consumers. An + unused extra-output gradient may be ``None`` and is treated as zero. +- Set ``output_to_caller=False`` for a channel tensor that should remain + internal. Removing a channel binding with ``channel=None`` restores that + output as public. +- Channel bindings are captured when an ``OperationFuser`` is + constructed, which locks them on every covered basic op. That includes: + + - constructing an ``OperationFuser`` or ``Sequential`` explicitly + - calling an op directly (``op(x)``), or a ``FusedOperation``, because + those paths build a transient ``OperationFuser([self])`` + + Later ``set_extra_input_channel`` / ``set_extra_output_channel`` calls + raise an error, so different routing requires constructing new + operations. Bind channels before the first forward call if the op will + later participate in a multi-op fuser. + +Channel-connected basic operations may still be replaced by registered +``FusedOperation`` implementations. If a fused operation contains both +the producer and consumer of a channel, its ``fuser_forward`` and +``fuser_backward`` implementations are responsible for routing the +tensor and its gradient between those basic operations. For a non-public +channel fully owned by one forward fusion, ``fuser_forward`` may return +``None`` in the corresponding basic-operation output slot. A tensor is +still required when the output is public or when a consumer is outside +that forward fusion. + Developer guide --------------- diff --git a/docs/examples/te_jax_integration.rst b/docs/examples/te_jax_integration.rst index a15a10e0b3..0d341cd9bc 100644 --- a/docs/examples/te_jax_integration.rst +++ b/docs/examples/te_jax_integration.rst @@ -28,8 +28,8 @@ Pick a topic - *Coming soon* - * - `Attention `_ - - *Coming soon* - - + - **Available** + - Single-GPU and context-parallel attention tutorials * - `Expert Parallelism `_ - *Coming soon* - diff --git a/examples/jax/encoder/requirements.txt b/examples/jax/encoder/requirements.txt index 141d35784a..dd0eb7e95c 100644 --- a/examples/jax/encoder/requirements.txt +++ b/examples/jax/encoder/requirements.txt @@ -1,4 +1,4 @@ datasets<4.0.0 flax>=0.7.1 -nltk>=3.8.2 +nltk>=3.8.2,<3.10.1 optax diff --git a/examples/jax/ep/bench/ep_bench.py b/examples/jax/ep/bench/ep_bench.py index 27ad8ca146..50b95e34f3 100644 --- a/examples/jax/ep/bench/ep_bench.py +++ b/examples/jax/ep/bench/ep_bench.py @@ -148,7 +148,7 @@ def main(): @jax.jit def run_prepare(idx): - tc, hm = tex_ep.ep_prepare(cfg, idx) + tc, _trt, hm = tex_ep.ep_prepare(cfg, idx) return tc, hm @jax.jit @@ -160,7 +160,7 @@ def run_dispatch(hm, idx, toks, w): @jax.jit def run_dispatch_vjp(idx, toks, w): - recv_t, recv_w, _hm, _tc = ep_dispatch(cfg, idx, toks, w, recv_capacity_per_rank) + recv_t, recv_w, _hm, _tc, _trt = ep_dispatch(cfg, idx, toks, w, recv_capacity_per_rank) recv_t = jax.lax.with_sharding_constraint(recv_t, NamedSharding(mesh, ep_spec_3d)) recv_w = jax.lax.with_sharding_constraint(recv_w, NamedSharding(mesh, ep_spec_2d)) return recv_t, recv_w diff --git a/examples/jax/ep/ep_moe.py b/examples/jax/ep/ep_moe.py index 671150d655..adf0f2ca0a 100644 --- a/examples/jax/ep/ep_moe.py +++ b/examples/jax/ep/ep_moe.py @@ -231,7 +231,7 @@ def _moe_layer(args, cfg, mesh, topk_idx, tokens, topk_w, local_kernels): local_kernels = jax.lax.with_sharding_constraint( local_kernels, NamedSharding(mesh, kernel_spec) ) - recv_tokens, recv_topk_w, handle_mem, _tc = ep_dispatch( + recv_tokens, recv_topk_w, handle_mem, _tc, _trt = ep_dispatch( cfg, topk_idx, tokens, topk_w, args.recv_capacity_per_rank ) recv_tokens = jax.lax.with_sharding_constraint(recv_tokens, NamedSharding(mesh, ep3)) diff --git a/examples/jax/ep/run_test_ep.sh b/examples/jax/ep/run_test_ep.sh index 1305ca6fd1..86aa6ca087 100755 --- a/examples/jax/ep/run_test_ep.sh +++ b/examples/jax/ep/run_test_ep.sh @@ -32,8 +32,8 @@ export PYTHONPATH="${TE_PATH}${PYTHONPATH:+:${PYTHONPATH}}" COORD="${COORD:-127.0.0.1:12345}" TEST_TIMEOUT_S="${TEST_TIMEOUT_S:-300}" -# Editable installs don't embed rpath; libtransformer_engine.so needs -# libnccl_ep.so.0 from the TE editable location at dlopen time. +# Editable installs don't embed rpath; the TE JAX extension needs +# libtransformer_engine.so from the TE editable location at dlopen time. TE_LIB_PATH=$(pip3 show transformer-engine 2>/dev/null \ | grep -E "Location:|Editable project location:" \ | tail -n 1 | awk '{print $NF}') diff --git a/examples/pytorch/ep/bench/ep_bench.py b/examples/pytorch/ep/bench/ep_bench.py index 2b7a2c62e5..f80a57015c 100644 --- a/examples/pytorch/ep/bench/ep_bench.py +++ b/examples/pytorch/ep/bench/ep_bench.py @@ -181,8 +181,9 @@ def main(): ep_group, num_experts=E, max_tokens_per_rank=T, - recv_capacity_per_rank=recv_pr, hidden_dim=H, + num_topk=K, + recv_capacity_per_rank=recv_pr, max_num_sms=args.max_num_sms, ) @@ -207,8 +208,6 @@ def main(): recv_capacity_per_rank=recv_pr, hidden_dim=H, num_local_experts=num_local_experts, - dispatch_recv_tokens=caller_recv_tokens, - combine_grad_expert_out=caller_grad_expert_out, ) tokens = tokens_hbm @@ -235,8 +234,13 @@ def main(): eo_p = recv_tokens.detach().clone().requires_grad_(True) # Stand-in callables; the cuda-graph branch below swaps in graphed versions. - fwd_bwd_dispatch_fn = lambda x: ep_dispatch(buffer, x, topk_idx, topk_w)[0] # noqa: E731 - fwd_bwd_combine_fn = lambda expert_out: ep_combine(buffer, expert_out) # noqa: E731 + # caller_* are None unless the caller opted in, matching ep_dispatch/ep_combine defaults. + fwd_bwd_dispatch_fn = lambda x: ep_dispatch( # noqa: E731 + buffer, x, topk_idx, topk_w, recv_tokens=caller_recv_tokens + )[0] + fwd_bwd_combine_fn = lambda expert_out: ep_combine( # noqa: E731 + buffer, expert_out, grad_out=caller_grad_expert_out + ) def _dispatch_raw(): _ep_dispatch_raw(buffer, topk_idx, tokens, topk_w, recv_tokens, recv_w) @@ -246,7 +250,7 @@ def _combine_raw(): _ep_combine_raw(buffer, expert_out, out_buf) def _ep_dispatch_fwd(): - ep_dispatch(buffer, tokens.detach(), topk_idx, topk_w) + ep_dispatch(buffer, tokens.detach(), topk_idx, topk_w, recv_tokens=caller_recv_tokens) def _ep_dispatch_fwd_bwd(): tokens_p.grad = None @@ -288,11 +292,11 @@ def _ep_combine_fwd_bwd(): # Graph fwd+bwd of the autograd-wrapped ops via make_graphed_callables. class _DispatchMod(torch.nn.Module): def forward(self, x): - return ep_dispatch(buffer, x, topk_idx, topk_w)[0] + return ep_dispatch(buffer, x, topk_idx, topk_w, recv_tokens=caller_recv_tokens)[0] class _CombineMod(torch.nn.Module): def forward(self, expert_out): - return ep_combine(buffer, expert_out) + return ep_combine(buffer, expert_out, grad_out=caller_grad_expert_out) disp_mod = _DispatchMod().cuda() comb_mod = _CombineMod().cuda() diff --git a/examples/pytorch/ep/bench/run_ep_bench.sh b/examples/pytorch/ep/bench/run_ep_bench.sh index fefecd7fa9..3b0977e4c3 100755 --- a/examples/pytorch/ep/bench/run_ep_bench.sh +++ b/examples/pytorch/ep/bench/run_ep_bench.sh @@ -26,10 +26,8 @@ if [ "${NSYS}" -eq 1 ] && [ "${KINETO}" -eq 1 ]; then fi SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)" RESULTS="${SCRIPT_DIR}/results" mkdir -p "${RESULTS}" -export PYTHONPATH="${TE_REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) NUM_GPUS="${NUM_GPUS:-${DETECTED_GPUS}}" diff --git a/examples/pytorch/ep/bench/run_nccl_ep_bench.sh b/examples/pytorch/ep/bench/run_nccl_ep_bench.sh index 8f6da04a00..ac5fcacc25 100755 --- a/examples/pytorch/ep/bench/run_nccl_ep_bench.sh +++ b/examples/pytorch/ep/bench/run_nccl_ep_bench.sh @@ -23,8 +23,8 @@ TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)" RESULTS="${SCRIPT_DIR}/results" mkdir -p "${RESULTS}" -BIN="${TE_REPO_ROOT}/3rdparty/nccl/build/test/nccl_ep/ep_bench" -LIB="${TE_REPO_ROOT}/3rdparty/nccl/build/lib" +BIN="${TE_REPO_ROOT}/3rdparty/nccl-extensions/build/test/nccl_ep/ep_bench" +LIB="${TE_REPO_ROOT}/3rdparty/nccl-extensions/build/lib" [ -x "${BIN}" ] || { echo "ep_bench not built at ${BIN}" >&2; exit 2; } NUM_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) diff --git a/examples/pytorch/ep/ep_moe.py b/examples/pytorch/ep/ep_moe.py index 149185f251..b47c334225 100644 --- a/examples/pytorch/ep/ep_moe.py +++ b/examples/pytorch/ep/ep_moe.py @@ -131,8 +131,9 @@ def main(): ep_group, num_experts=num_experts, max_tokens_per_rank=T, - recv_capacity_per_rank=recv_pr, hidden_dim=args.hidden, + num_topk=args.top_k, + recv_capacity_per_rank=recv_pr, ) try: _run_layer( @@ -182,15 +183,13 @@ def _run_layer(args, rank, world_size, ep_size, num_experts, num_local_experts, recv_capacity_per_rank=recv_pr, hidden_dim=args.hidden, num_local_experts=num_local_experts, - dispatch_recv_tokens=recv_tokens, - combine_grad_expert_out=grad_expert_out, ) - recv_t, recv_w_out, _tc = ep_dispatch(buffer, tokens, topk_idx, topk_w) + recv_t, recv_w_out, _tc = ep_dispatch(buffer, tokens, topk_idx, topk_w, recv_tokens=recv_tokens) expert_out = _batched_expert_linear(recv_t, kernels_local, num_local_experts) # Apply per-slot topk weighting before combine. expert_out = expert_out * recv_w_out.unsqueeze(-1).to(expert_out.dtype) - out = ep_combine(buffer, expert_out) + out = ep_combine(buffer, expert_out, grad_out=grad_expert_out) loss = 0.5 * (out.float() ** 2).sum() loss.backward() diff --git a/examples/pytorch/ep/run_test_ep.sh b/examples/pytorch/ep/run_test_ep.sh index 13b41f4cb2..d8e6b50556 100755 --- a/examples/pytorch/ep/run_test_ep.sh +++ b/examples/pytorch/ep/run_test_ep.sh @@ -17,7 +17,6 @@ if [ "${NUM_GPUS}" -gt 8 ]; then NUM_GPUS=8; fi : ${TEST_TIMEOUT_S:=120} SCRIPT="${TE_PATH}/examples/pytorch/ep/ep_moe.py" -export PYTHONPATH="${TE_PATH}${PYTHONPATH:+:${PYTHONPATH}}" # Stage JIT cubins on tmpfs for fast iteration. : ${NCCL_EP_JIT_CACHE_DIR:="${TMPDIR:-/tmp}/nccl_ep_jit_cache_$(id -u)"} diff --git a/qa/L0_jax_unittest/test.sh b/qa/L0_jax_unittest/test.sh index 5d833792cc..51ff46ba79 100644 --- a/qa/L0_jax_unittest/test.sh +++ b/qa/L0_jax_unittest/test.sh @@ -21,8 +21,8 @@ FAILED_CASES="" export NVTE_JAX_TEST_TIMING=1 -pip3 install "nltk>=3.8.2" || error_exit "Failed to install nltk" -pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" +pip3 install "nltk>=3.8.2,<3.10.1" || error_exit "Failed to install nltk" +pip3 install pytest==8.2.1 pytest-timeout==2.4.0 || error_exit "Failed to install pytest dependencies" : ${TE_PATH:=/opt/transformerengine} : ${XML_LOG_DIR:=/logs} @@ -40,7 +40,6 @@ pip3 install -r $TE_PATH/examples/jax/encoder/requirements.txt || error_exit "Fa export XLA_FLAGS="${XLA_FLAGS} --xla_gpu_deterministic_ops" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_test_single_gpu_encoder.xml $TE_PATH/examples/jax/encoder/test_single_gpu_encoder.py || test_fail "test_single_gpu_encoder.py" # Test without custom calls -export XLA_FLAGS="${XLA_FLAGS} --xla_gpu_deterministic_ops" NVTE_JAX_CUSTOM_CALLS="false" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_test_single_gpu_encoder_without_custom_call.xml $TE_PATH/examples/jax/encoder/test_single_gpu_encoder.py || test_fail "test_single_gpu_encoder.py without custom calls" # Exercise the docs/examples/jax tutorials. The multi-GPU tests are diff --git a/qa/L0_pytorch_debug_unittest/test.sh b/qa/L0_pytorch_debug_unittest/test.sh index 3efa462628..36efe485f5 100644 --- a/qa/L0_pytorch_debug_unittest/test.sh +++ b/qa/L0_pytorch_debug_unittest/test.sh @@ -22,6 +22,11 @@ FAILED_CASES="" : ${XML_LOG_DIR:=/logs} mkdir -p "$XML_LOG_DIR" +# L0 keeps one mature FlashAttention generation; L3 owns newer-generation coverage. +export NVTE_FLASH_ATTN_V2=1 +export NVTE_FLASH_ATTN_V3=0 +export NVTE_FLASH_ATTN_V4=0 + # Config with the dummy feature which prevents nvinspect from being disabled. # Nvinspect will be disabled if no feature is active. : ${NVTE_TEST_NVINSPECT_DUMMY_CONFIG_FILE:=$TE_PATH/tests/pytorch/debug/test_configs/dummy_feature.yaml} diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 91d3be63d0..14a5f4fe3d 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -22,6 +22,11 @@ set -x : ${XML_LOG_DIR:=/logs} mkdir -p "$XML_LOG_DIR" +# L0 keeps one mature FlashAttention generation; L3 owns newer-generation coverage. +export NVTE_FLASH_ATTN_V2=1 +export NVTE_FLASH_ATTN_V3=0 +export NVTE_FLASH_ATTN_V4=0 + pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/test_sanity.py || test_fail "test_sanity.py" @@ -40,21 +45,29 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_torch_compile.xm python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8blockwisetensor.xml $TE_PATH/tests/pytorch/test_float8blockwisetensor.py || test_fail "test_float8blockwisetensor.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_blockwise_scaling_exact.xml $TE_PATH/tests/pytorch/test_float8_blockwise_scaling_exact.py || test_fail "test_float8_blockwise_scaling_exact.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_blockwise_gemm_exact.xml $TE_PATH/tests/pytorch/test_float8_blockwise_gemm_exact.py || test_fail "test_float8_blockwise_gemm_exact.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_current_scaling_exact.xml $TE_PATH/tests/pytorch/test_float8_current_scaling_exact.py || test_fail "test_float8_current_scaling_exact.py" NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/test_grouped_tensor.xml $TE_PATH/tests/pytorch/test_grouped_tensor.py || test_fail "test_grouped_tensor.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_gqa.xml $TE_PATH/tests/pytorch/test_gqa.py || test_fail "test_gqa.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_qk_norm.xml $TE_PATH/tests/pytorch/test_qk_norm.py || test_fail "test_qk_norm.py" 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" +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" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_permutation.xml $TE_PATH/tests/pytorch/test_permutation.py || test_fail "test_permutation.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_parallel_cross_entropy.xml $TE_PATH/tests/pytorch/test_parallel_cross_entropy.py || test_fail "test_parallel_cross_entropy.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading.xml $TE_PATH/tests/pytorch/test_cpu_offloading.py || test_fail "test_cpu_offloading.py" NVTE_FLASH_ATTN=0 NVTE_CPU_OFFLOAD_V1=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading_v1.xml $TE_PATH/tests/pytorch/test_cpu_offloading_v1.py || test_fail "test_cpu_offloading_v1.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hybrid_quantization.xml $TE_PATH/tests/pytorch/test_hybrid_quantization.py || test_fail "test_hybrid_quantization.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_identity_quantizer.xml $TE_PATH/tests/pytorch/test_identity_quantizer.py || test_fail "test_identity_quantizer.py" NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "test_attention.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_flex_attention.xml $TE_PATH/tests/pytorch/attention/test_flex_attention.py || test_fail "test_flex_attention.py" NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_linear_mxfp8_attention.xml $TE_PATH/tests/pytorch/attention/test_linear_mxfp8_attention.py || test_fail "test_linear_mxfp8_attention.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_mla_q_uproj.xml $TE_PATH/tests/pytorch/attention/test_fused_mla_q_uproj.py || test_fail "test_fused_mla_q_uproj.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_kv_cache.xml $TE_PATH/tests/pytorch/attention/test_kv_cache.py || test_fail "test_kv_cache.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cu_seqlens_cache.xml $TE_PATH/tests/pytorch/attention/test_cu_seqlens_cache.py || test_fail "test_cu_seqlens_cache.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hf_integration.xml $TE_PATH/tests/pytorch/test_hf_integration.py || test_fail "test_hf_integration.py" export NVTE_TEST_CHECKPOINT_ARTIFACT_PATH=$TE_PATH/artifacts/tests/pytorch/test_checkpoint if [ ! -d "$NVTE_TEST_CHECKPOINT_ARTIFACT_PATH" ]; then @@ -63,9 +76,12 @@ fi python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_checkpoint.xml $TE_PATH/tests/pytorch/test_checkpoint.py || test_fail "test_checkpoint.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_router.xml $TE_PATH/tests/pytorch/test_fused_router.py || test_fail "test_fused_router.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_partial_cast.xml $TE_PATH/tests/pytorch/test_partial_cast.py || test_fail "test_partial_cast.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mxfp8_2d_quantize.xml $TE_PATH/tests/pytorch/test_mxfp8_2d_quantize.py || test_fail "test_mxfp8_2d_quantize.py" # Disable autotuning to make unittests faster. In addition, disable TF32 path to fully align with the pytorch reference implementation's precision NVTE_DISABLE_TRITON_AUTOTUNING=1 NVIDIA_TF32_OVERRIDE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mhc.xml $TE_PATH/tests/pytorch/test_mhc.py || test_fail "test_mhc.py" +NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_DISABLE_TRITON_AUTOTUNING=1 NVIDIA_TF32_OVERRIDE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mhc_deterministic.xml $TE_PATH/tests/pytorch/test_mhc.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_mhc.py" PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_grouped_linear.xml $TE_PATH/tests/pytorch/test_grouped_linear.py || test_fail "test_grouped_linear.py" +PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_ops_grouped_linear_distributed_weight.xml $TE_PATH/tests/pytorch/test_ops_grouped_linear_distributed_weight.py || test_fail "test_ops_grouped_linear_distributed_weight.py" NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_grouped_mlp.xml $TE_PATH/tests/pytorch/test_grouped_mlp.py || test_fail "test_grouped_mlp.py" if [ "$RET" -ne 0 ]; then diff --git a/qa/L1_jax_distributed_unittest/test.sh b/qa/L1_jax_distributed_unittest/test.sh index 8e0ef2c267..6feec1406b 100644 --- a/qa/L1_jax_distributed_unittest/test.sh +++ b/qa/L1_jax_distributed_unittest/test.sh @@ -31,12 +31,12 @@ python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/py python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_mlp.xml $TE_PATH/tests/jax/test_distributed_layernorm_mlp.py || test_fail "test_distributed_layernorm_mlp.py" -# XLA_FLAGS to WAR for test_distributed_softmax issue with NCCL -# TODO(Kshitij): remove when NCCL issue is fixed -XLA_FLAGS="$XLA_FLAGS --xla_gpu_enable_nccl_comm_splitting=false" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_softmax.xml $TE_PATH/tests/jax/test_distributed_softmax.py || test_fail "test_distributed_softmax.py" - python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_fused_attn.xml $TE_PATH/tests/jax/test_distributed_fused_attn.py || test_fail "test_distributed_fused_attn.py" +# XLA_FLAGS to WAR for test_distributed_softmax issues with NCCL and generic async all-reduce. +# TODO(KshitijLakhani): remove when NCCL issues and openxla/xla#46938 are fixed. +XLA_FLAGS="$XLA_FLAGS --xla_gpu_enable_nccl_comm_splitting=false --xla_gpu_disable_async_collectives=ALLREDUCE" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_softmax.xml $TE_PATH/tests/jax/test_distributed_softmax.py || test_fail "test_distributed_softmax.py" + # NCCL EP multi-process suite. Self-skips on <4 GPUs. TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/multi_process_launch_ep.sh || test_fail "test_multi_process_ep.py" diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh index 50a51353d1..e0c92849f8 100644 --- a/qa/L1_pytorch_distributed_unittest/test.sh +++ b/qa/L1_pytorch_distributed_unittest/test.sh @@ -18,12 +18,15 @@ FAILED_CASES="" : ${TE_PATH:=/opt/transformerengine} : ${XML_LOG_DIR:=/logs} +# FA4 coverage belongs to the dedicated attention-backend suites. +export NVTE_FLASH_ATTN_V4=0 mkdir -p "$XML_LOG_DIR" pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" # Run CP tests (deterministic + non-deterministic) first so they can be parallelized. # Each needs 4 GPUs, so >=8 GPUs allows them to run concurrently on disjoint GPU sets. +# Main's CP implementation supports FA2/FA3. NUM_GPUS=$(python3 -c "import torch; print(torch.cuda.device_count())") echo "Detected $NUM_GPUS GPU(s)" if [ "$NUM_GPUS" -ge 8 ]; then @@ -44,10 +47,11 @@ python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/ python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics.xml $TE_PATH/tests/pytorch/distributed/test_numerics.py || test_fail "test_numerics.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics_exact.xml $TE_PATH/tests/pytorch/distributed/test_numerics_exact.py || test_fail "test_numerics_exact.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops.xml $TE_PATH/tests/pytorch/distributed/test_fusible_ops.py || test_fail "test_fusible_ops.py" -python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_torch_fsdp2.xml $TE_PATH/tests/pytorch/distributed/test_torch_fsdp2.py || test_fail "test_torch_fsdp2.py" +python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_torch_fsdp2.xml $TE_PATH/tests/pytorch/distributed/test_torch_fsdp2.py -k "not hybrid" || test_fail "test_torch_fsdp2.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_comm_gemm_overlap.xml $TE_PATH/tests/pytorch/distributed/test_comm_gemm_overlap.py || test_fail "test_comm_gemm_overlap.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops_with_userbuffers.xml $TE_PATH/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py || test_fail "test_fusible_ops_with_userbuffers.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_cp_utils.xml $TE_PATH/tests/pytorch/attention/test_cp_utils.py || test_fail "test_cp_utils.py" +python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_cu_seqlens_cache.xml $TE_PATH/tests/pytorch/attention/test_cu_seqlens_cache.py || test_fail "test_cu_seqlens_cache.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_cast_master_weights_to_fp8.xml $TE_PATH/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py || test_fail "test_cast_master_weights_to_fp8.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_newton_schulz.xml $TE_PATH/tests/pytorch/distributed/test_newton_schulz.py || test_fail "test_newton_schulz.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_ep.xml $TE_PATH/tests/pytorch/distributed/test_ep.py || test_fail "test_ep.py" diff --git a/qa/L1_pytorch_hybrid_distributed_unittest/test.sh b/qa/L1_pytorch_hybrid_distributed_unittest/test.sh new file mode 100644 index 0000000000..a2c06266b9 --- /dev/null +++ b/qa/L1_pytorch_hybrid_distributed_unittest/test.sh @@ -0,0 +1,33 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +function error_exit() { + echo "Error: $1" + exit 1 +} + +function test_fail() { + RET=1 + FAILED_CASES="$FAILED_CASES $1" + echo "Error: sub-test failed: $1" +} + +RET=0 +FAILED_CASES="" + +: ${TE_PATH:=/opt/transformerengine} +: ${XML_LOG_DIR:=/logs} +mkdir -p "$XML_LOG_DIR" + +pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" + +python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_torch_fsdp2_hybrid.xml $TE_PATH/tests/pytorch/distributed/test_torch_fsdp2.py -k "hybrid" || test_fail "hybrid test_torch_fsdp2.py" +python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_hybrid_tp_sp.xml $TE_PATH/tests/pytorch/distributed/test_hybrid_tp_sp.py || test_fail "test_hybrid_tp_sp.py" + +if [ "$RET" -ne 0 ]; then + echo "Error in the following test cases:$FAILED_CASES" + exit 1 +fi +echo "All tests passed" +exit 0 diff --git a/qa/L2_jax_distributed_unittest/test.sh b/qa/L2_jax_distributed_unittest/test.sh index 330b254e7d..3dcd85ede2 100644 --- a/qa/L2_jax_distributed_unittest/test.sh +++ b/qa/L2_jax_distributed_unittest/test.sh @@ -12,7 +12,18 @@ export NVTE_JAX_TEST_TIMING=1 mkdir -p "$XML_LOG_DIR" # Use --xla_gpu_enable_triton_gemm=false to ensure the reference JAX implementation we are using is accurate. -XLA_FLAGS="$XLA_FLAGS --xla_gpu_enable_triton_gemm=false" NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_* +common_xla_flags="$XLA_FLAGS --xla_gpu_enable_triton_gemm=false" + +distributed_tests=() +for test_file in $TE_PATH/tests/jax/test_distributed_*.py; do + if [ "$test_file" != "$TE_PATH/tests/jax/test_distributed_softmax.py" ]; then + distributed_tests+=("$test_file") + fi +done +XLA_FLAGS="$common_xla_flags" NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml "${distributed_tests[@]}" + +# Work around the XLA 26.08 generic async all-reduce deadlock (openxla/xla#46938). +XLA_FLAGS="$common_xla_flags --xla_gpu_enable_nccl_comm_splitting=false --xla_gpu_disable_async_collectives=ALLREDUCE" NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_softmax.xml $TE_PATH/tests/jax/test_distributed_softmax.py # NCCL EP multi-process suite. The launcher skips when fewer than 4 GPUs or no NVLink is detected. TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/multi_process_launch_ep.sh diff --git a/qa/L2_jax_unittest/test.sh b/qa/L2_jax_unittest/test.sh index f455ec0df3..194b9672a1 100644 --- a/qa/L2_jax_unittest/test.sh +++ b/qa/L2_jax_unittest/test.sh @@ -21,8 +21,8 @@ FAILED_CASES="" export NVTE_JAX_TEST_TIMING=1 -pip3 install "nltk>=3.8.2" || error_exit "Failed to install nltk" -pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" +pip3 install "nltk>=3.8.2,<3.10.1" || error_exit "Failed to install nltk" +pip3 install pytest==8.2.1 pytest-timeout==2.4.0 || error_exit "Failed to install pytest dependencies" : ${TE_PATH:=/opt/transformerengine} : ${XML_LOG_DIR:=/logs} diff --git a/qa/L3_pytorch_FA_versions_test/test.sh b/qa/L3_pytorch_FA_versions_test/test.sh index 30f1fc38c0..047358b301 100644 --- a/qa/L3_pytorch_FA_versions_test/test.sh +++ b/qa/L3_pytorch_FA_versions_test/test.sh @@ -25,40 +25,63 @@ pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" # Limit parallel build jobs to avoid overwhelming system resources export MAX_JOBS=32 +# Checkpoint for FP8 delayed scaling uses pickle +export NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 + # Iterate over Flash Attention versions sm_arch=`python3 -c "import torch; sm = torch.cuda.get_device_capability(0); print(sm[0]*10+sm[1])"` export FLASH_ATTN_CUDA_ARCHS=$sm_arch -# CP tests are expensive and run only once per arch: -# - sm90 (H100): FA3 (3.0.0b1) - context_parallel.py only supports FA3 on Hopper -# - sm>90 (B200): latest FA4 - FA3 is not built/installed for sm>90 -# Non-CP tests still run for every FA version in the array. +# Run one architecture-owned FlashAttention generation. CP remains FA3-only +# until the production selector and runner support FA4 CP end to end. +CP_FA_VERSION="" if [ $sm_arch -gt 90 ] then - FA_versions=(2.8.3 4.0.0b11) - CP_FA_VERSION="${FA_versions[-1]}" + FA_versions=(4.0.0b11) elif [ $sm_arch -eq 90 ] then - FA_versions=(2.8.3 3.0.0b1 4.0.0b11) + FA_versions=(3.0.0b1) CP_FA_VERSION="3.0.0b1" +else + error_exit "No L3 FlashAttention generation is defined for sm${sm_arch}" fi for fa_version in "${FA_versions[@]}" do + # The FA distributions share the flash_attn namespace. Keep exactly one + # installed so import-time discovery and the iteration label cannot disagree. + pip3 uninstall -y flash-attn flash-attn-3 flash-attn-4 \ + || error_exit "Failed to isolate Flash Attention $fa_version" + export NVTE_FLASH_ATTN_V2=0 + export NVTE_FLASH_ATTN_V3=0 + export NVTE_FLASH_ATTN_V4=0 + # Build Flash Attention if [ "${fa_version}" \< "3.0.0" ] then - pip3 install flash-attn==${fa_version} --no-build-isolation + export NVTE_FLASH_ATTN_V2=1 + pip3 install flash-attn==${fa_version} --no-build-isolation \ + || error_exit "Failed to install Flash Attention $fa_version" elif [[ "${fa_version}" == 4.* ]] then - pip3 install flash-attn-4==${fa_version} nvidia-cutlass-dsl[cu13]==4.4.2 --no-build-isolation + export NVTE_FLASH_ATTN_V4=1 + # FA4 is intentionally last in every version array. Its b11 test pin needs + # CUTLASS DSL 4.4.2, so replace the image-matched stack only for this final + # iteration; later iterations would otherwise need that stack restored. + pip3 uninstall -y nvidia-cutlass-dsl nvidia-cutlass-dsl-libs-base \ + nvidia-cutlass-dsl-libs-cu12 nvidia-cutlass-dsl-libs-cu13 \ + || error_exit "Failed to isolate CUTLASS DSL for Flash Attention $fa_version" + pip3 install flash-attn-4==${fa_version} nvidia-cutlass-dsl[cu13]==4.4.2 \ + --no-build-isolation || error_exit "Failed to install Flash Attention $fa_version" else + export NVTE_FLASH_ATTN_V3=1 # FA3 source build (~20 min). Skip if FA3 is already installed. if python3 -c "import flash_attn_3" 2>/dev/null; then echo "FA3 already installed (from base image); skipping source build" else git clone https://github.com/Dao-AILab/flash-attention.git - cd flash-attention/hopper && python setup.py install + cd flash-attention/hopper && python setup.py install \ + || error_exit "Failed to install Flash Attention $fa_version" cd ../../ fi fi @@ -77,14 +100,16 @@ do XML_ATTN="$XML_LOG_DIR/pytest_test_attention_fa${fa_tag}.xml" XML_CP="$XML_LOG_DIR/pytest_test_attention_with_cp_fa${fa_tag}.xml" - if [ "$fa_version" = "$CP_FA_VERSION" ]; then + # test_attention.py reloads its own trusted delayed-scaling FP8 checkpoint, + # whose legacy extra state requires an explicit pickle opt-in. + if [ -n "$CP_FA_VERSION" ] && [ "$fa_version" = "$CP_FA_VERSION" ]; then echo "Running CP tests with FA $fa_version (CP version for sm$sm_arch)" if [ "$NUM_GPUS" -ge 5 ]; then CP_NUM_GPUS=$(( NUM_GPUS - 1 > 4 ? 4 : NUM_GPUS - 1 )) CP_GPUS=$(seq -s, 1 $CP_NUM_GPUS) echo "Running tests in parallel: test_attention.py on GPU 0, test_attention_with_cp.py on GPUs $CP_GPUS ($CP_NUM_GPUS GPUs)" - CUDA_VISIBLE_DEVICES=0 NVTE_TORCH_COMPILE=0 python3 -m pytest -v -s \ + CUDA_VISIBLE_DEVICES=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 python3 -m pytest -v -s \ --junitxml=$XML_ATTN \ $TE_PATH/tests/pytorch/attention/test_attention.py & PID_ATTN=$! @@ -98,12 +123,16 @@ do wait $PID_CP || test_fail "test_attention_with_cp.py (FA $fa_version)" else echo "Running tests sequentially: need >=5 GPUs for parallel execution (1 for test_attention + 4 for test_attention_with_cp)" - NVTE_TORCH_COMPILE=0 python3 -m pytest -v -s --junitxml=$XML_ATTN $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "test_attention.py (FA $fa_version)" + NVTE_TORCH_COMPILE=0 NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 python3 -m pytest -v -s --junitxml=$XML_ATTN $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "test_attention.py (FA $fa_version)" NVTE_TORCH_COMPILE=0 python3 -m pytest -v -s --junitxml=$XML_CP $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "test_attention_with_cp.py (FA $fa_version)" fi else - echo "Skipping CP tests for FA $fa_version (CP only runs with FA $CP_FA_VERSION on sm$sm_arch)" - NVTE_TORCH_COMPILE=0 python3 -m pytest -v -s --junitxml=$XML_ATTN $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "test_attention.py (FA $fa_version)" + if [ -n "$CP_FA_VERSION" ]; then + echo "Skipping CP tests for FA $fa_version (CP uses FA $CP_FA_VERSION on sm$sm_arch)" + else + echo "CP tests are not scheduled for the FA generation on sm$sm_arch" + fi + NVTE_TORCH_COMPILE=0 NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 python3 -m pytest -v -s --junitxml=$XML_ATTN $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "test_attention.py (FA $fa_version)" fi done diff --git a/setup.py b/setup.py index 02c00fecd9..1a84fe3546 100644 --- a/setup.py +++ b/setup.py @@ -34,6 +34,7 @@ remove_dups, min_python_version_str, nccl_ep_enabled, + get_max_jobs_for_parallel_build, ) frameworks = get_frameworks() @@ -288,11 +289,11 @@ def _discover_nccl_home() -> str: def build_nccl_ep_submodule() -> str: - """Build libnccl_ep.a from the 3rdparty/nccl submodule and return NCCL_HOME.""" - nccl_root = current_file_path / "3rdparty" / "nccl" - if not (nccl_root / "Makefile").exists(): + """Build libnccl_ep.a from the 3rdparty/nccl-extensions submodule and return NCCL_HOME.""" + nccl_root = current_file_path / "3rdparty" / "nccl-extensions" + if not (nccl_root / "nccl_ep" / "Makefile").exists(): raise RuntimeError( - f"NCCL submodule not found at {nccl_root}. " + f"NCCL EP submodule not found at {nccl_root}. " "Run `git submodule update --init --recursive`." ) @@ -331,7 +332,7 @@ def build_nccl_ep_submodule() -> str: ) gencode = " ".join(f"-gencode=arch=compute_{a},code=sm_{a}" for a in arch_list) - nproc = os.cpu_count() or 8 + nproc = get_max_jobs_for_parallel_build() env = os.environ.copy() env["NVCC_GENCODE"] = gencode # NCCL EP needs the core NCCL headers + libnccl.so; write NCCL EP build @@ -348,13 +349,14 @@ def build_nccl_ep_submodule() -> str: "rebuilding libnccl_ep.a" ) subprocess.check_call( - ["make", "-C", "contrib/nccl_ep", "clean"], + ["make", "-C", "nccl_ep", "clean"], cwd=str(nccl_root), env=env, ) print(f"[NCCL EP] Building libnccl_ep.a (gencode='{gencode}')") + make_jobs = f"-j{nproc}" if nproc else "-j" subprocess.check_call( - ["make", "-j", str(nproc), "-C", "contrib/nccl_ep", "lib"], + ["make", make_jobs, "-C", "nccl_ep", "lib"], cwd=str(nccl_root), env=env, ) diff --git a/tests/cpp/operator/test_cast_float8blockwise_grouped.cu b/tests/cpp/operator/test_cast_float8blockwise_grouped.cu index bc9f104e17..90418d6287 100644 --- a/tests/cpp/operator/test_cast_float8blockwise_grouped.cu +++ b/tests/cpp/operator/test_cast_float8blockwise_grouped.cu @@ -368,6 +368,7 @@ struct TestConfig { ScalingDir dir; std::vector first_dims; size_t K; + bool force_pow_2_scales; }; class GroupedFP8BlockwiseTestSuite : public ::testing::TestWithParam {}; @@ -375,7 +376,7 @@ class GroupedFP8BlockwiseTestSuite : public ::testing::TestWithParam TEST_P(GroupedFP8BlockwiseTestSuite, Test) { const TestConfig& cfg = GetParam(); perform_test(cfg.shape_rep, cfg.block_dim, cfg.dir, cfg.first_dims, cfg.K, - /*force_pow_2_scales=*/false, /*epsilon=*/0.0f); + cfg.force_pow_2_scales, /*epsilon=*/0.0f); } std::vector make_configs() { @@ -387,11 +388,13 @@ std::vector make_configs() { for (auto bd : {BlockDim::ONE_D, BlockDim::TWO_D}) { for (auto dir : {ScalingDir::ROWWISE, ScalingDir::COLWISE, ScalingDir::BOTH}) { for (size_t K : Ks) { - for (const auto& v : uniform) { - configs.push_back({ShapeRep::SAME_BOTH_DIMS, bd, dir, v, K}); - } - for (const auto& v : jagged) { - configs.push_back({ShapeRep::VARYING_FIRST_DIM, bd, dir, v, K}); + for (bool pow2 : {false, true}) { + for (const auto& v : uniform) { + configs.push_back({ShapeRep::SAME_BOTH_DIMS, bd, dir, v, K, pow2}); + } + for (const auto& v : jagged) { + configs.push_back({ShapeRep::VARYING_FIRST_DIM, bd, dir, v, K, pow2}); + } } } } @@ -408,6 +411,7 @@ std::string make_name(const ::testing::TestParamInfo& info) { s += "_K" + std::to_string(c.K) + "_N" + std::to_string(c.first_dims.size()); s += "_M"; for (size_t m : c.first_dims) s += "_" + std::to_string(m); + s += (c.force_pow_2_scales ? "_POW2" : "_FP32SC"); return s; } diff --git a/tests/cpp/operator/test_cast_mxfp8.cu b/tests/cpp/operator/test_cast_mxfp8.cu index 738c25d7d4..5b99f2f8cd 100644 --- a/tests/cpp/operator/test_cast_mxfp8.cu +++ b/tests/cpp/operator/test_cast_mxfp8.cu @@ -6,13 +6,17 @@ * See LICENSE for license information. ************************************************************************/ +#include + #include #include #include +#include #include #include #include +#include #include "../test_common.h" #include "transformer_engine/transformer_engine.h" @@ -38,6 +42,12 @@ enum ActivationType { SReLU }; +enum MXFP82DScalingDirection { + RowwiseOnly, + ColwiseOnly, + Bidirectional +}; + template void compute_ref(const ProcessingMethod processing_method, float (*OP)(const float), @@ -170,6 +180,70 @@ void compute_ref(const ProcessingMethod processing_method, } } +template +void compute_ref_2d_quantize(const bool rowwise, + const bool colwise, + const InputType* input, + OutputType* output_rowwise, + OutputType* output_colwise, + fp8e8m0* output_scales_rowwise, + fp8e8m0* output_scales_colwise, + const size_t rows, + const size_t cols, + const size_t scales_stride_rowwise, + const size_t scales_stride_colwise) { + const size_t tile_size_Y = 32; + const size_t tile_size_X = 32; + const size_t tiles_num_Y = (rows + tile_size_Y - 1) / tile_size_Y; + const size_t tiles_num_X = (cols + tile_size_X - 1) / tile_size_X; + + #pragma omp parallel for collapse(2) proc_bind(spread) + for (size_t tile_Y = 0; tile_Y < tiles_num_Y; ++tile_Y) { + for (size_t tile_X = 0; tile_X < tiles_num_X; ++tile_X) { + const size_t i_min = tile_Y * tile_size_Y; + const size_t i_max = std::min(i_min + tile_size_Y, rows); + const size_t j_min = tile_X * tile_size_X; + const size_t j_max = std::min(j_min + tile_size_X, cols); + + float block_amax = 0.0f; + for (size_t i = i_min; i < i_max; ++i) { + for (size_t j = j_min; j < j_max; ++j) { + const size_t idx = i * cols + j; + block_amax = std::max(block_amax, std::abs(static_cast(input[idx]))); + } + } + + const fp8e8m0 biased_exponent = + float_to_e8m0(block_amax * Quantized_Limits::max_reciprocal()); + const float scale_reciprocal = exp2f_rcp(biased_exponent); + + if (rowwise) { + for (size_t i = i_min; i < i_max; ++i) { + output_scales_rowwise[i * scales_stride_rowwise + tile_X] = biased_exponent; + for (size_t j = j_min; j < j_max; ++j) { + const size_t idx = i * cols + j; + output_rowwise[idx] = + static_cast(static_cast(input[idx]) * + scale_reciprocal); + } + } + } + + if (colwise) { + for (size_t j = j_min; j < j_max; ++j) { + output_scales_colwise[tile_Y * scales_stride_colwise + j] = biased_exponent; + for (size_t i = i_min; i < i_max; ++i) { + const size_t idx = i * cols + j; + output_colwise[idx] = + static_cast(static_cast(input[idx]) * + scale_reciprocal); + } + } + } + } + } +} + /** * Scaling along single dimension (either rows or columns) * Produces one set of output data and the corresponding data of the fused operation (dbias): @@ -569,6 +643,106 @@ void performTest_x2(const ProcessingMethod processing_method, } } +template +void performTest_2d_quantize(const std::vector& shape, + const MXFP82DScalingDirection scaling_direction, + InputsFillCase fill_case) { + using namespace test; + using EncodingType = fp32; + DType itype = TypeInfo::dtype; + DType otype = TypeInfo::dtype; + + if (shape.size() < 2) { + GTEST_SKIP(); + } + + const size_t rows = first_dimension(shape); + const size_t cols = last_dimension(shape); + + const bool rowwise = scaling_direction != MXFP82DScalingDirection::ColwiseOnly; + const bool colwise = scaling_direction != MXFP82DScalingDirection::RowwiseOnly; + + const std::array scale_dims_rowwise = get_scale_tensor_dims(rows, cols, 1, 32); + const std::array scale_dims_colwise = get_scale_tensor_dims(rows, cols, 32, 1); + + const size_t unpadded_blocks_Y_rowwise = scale_dims_rowwise[0]; + const size_t unpadded_blocks_X_rowwise = scale_dims_rowwise[1]; + const size_t blocks_Y_rowwise = scale_dims_rowwise[2]; + const size_t blocks_X_rowwise = scale_dims_rowwise[3]; + const size_t scales_stride_rowwise = blocks_X_rowwise; + + const size_t unpadded_blocks_Y_colwise = scale_dims_colwise[0]; + const size_t unpadded_blocks_X_colwise = scale_dims_colwise[1]; + const size_t blocks_Y_colwise = scale_dims_colwise[2]; + const size_t blocks_X_colwise = scale_dims_colwise[3]; + const size_t scales_stride_colwise = blocks_X_colwise; + + Tensor input("input", shape, itype); + Tensor output("output", shape, otype, rowwise, colwise, NVTE_MXFP8_1D_SCALING); + + std::unique_ptr ref_output_rowwise = std::make_unique(rows * cols); + std::unique_ptr ref_output_colwise = std::make_unique(rows * cols); + std::unique_ptr ref_scales_rowwise = + std::make_unique(blocks_Y_rowwise * blocks_X_rowwise); + std::unique_ptr ref_scales_colwise = + std::make_unique(blocks_Y_colwise * blocks_X_colwise); + std::fill_n(ref_scales_rowwise.get(), blocks_Y_rowwise * blocks_X_rowwise, 0); + std::fill_n(ref_scales_colwise.get(), blocks_Y_colwise * blocks_X_colwise, 0); + + fillCase(&input, fill_case); + + QuantizationConfigWrapper quant_config; + quant_config.set_mxfp8_2d_quantization(true); + nvte_quantize_v2(input.data(), output.data(), quant_config, 0); + + cudaDeviceSynchronize(); + auto err = cudaGetLastError(); + ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); + + compute_ref_2d_quantize( + rowwise, + colwise, + input.rowwise_cpu_dptr(), + ref_output_rowwise.get(), + ref_output_colwise.get(), + ref_scales_rowwise.get(), + ref_scales_colwise.get(), + rows, + cols, + scales_stride_rowwise, + scales_stride_colwise); + + const size_t scale_diff_abs_tolerance = 0; + const double abs_tolerable_mismatches_limit = 0.0; + const double rel_tolerable_mismatches_limit = 0.0; + + auto [atol, rtol] = getTolerances(otype); + + if (rowwise) { + size_t mismatches_scales_rowwise = 0; + compare_scaling_factors("scales_rowwise", output.rowwise_cpu_scale_inv_ptr(), + ref_scales_rowwise.get(), unpadded_blocks_Y_rowwise, + unpadded_blocks_X_rowwise, scales_stride_rowwise, + mismatches_scales_rowwise, + scale_diff_abs_tolerance, + abs_tolerable_mismatches_limit, + rel_tolerable_mismatches_limit); + compareResults("output_rowwise", output, ref_output_rowwise.get(), true, atol, rtol, true); + } + + if (colwise) { + size_t mismatches_scales_colwise = 0; + compare_scaling_factors("scales_colwise", output.columnwise_cpu_scale_inv_ptr(), + ref_scales_colwise.get(), unpadded_blocks_Y_colwise, + unpadded_blocks_X_colwise, scales_stride_colwise, + mismatches_scales_colwise, + scale_diff_abs_tolerance, + abs_tolerable_mismatches_limit, + rel_tolerable_mismatches_limit); + compareResults("output_colwise", output, ref_output_colwise.get(), false, atol, rtol, true); + } +} + std::vector> matrix_sizes = { {1, 16}, {16, 48}, @@ -580,6 +754,18 @@ std::vector> matrix_sizes = { {8192, 7168}, }; +std::vector> matrix_sizes_2d_quantize = { + {1, 16}, + {16, 48}, + {65, 80}, + {127, 400}, + {128, 128}, + {993, 512}, + {8, 32, 1024}, + {16, 8, 4, 512}, + {8192, 7168}, +}; + std::vector> block_sizes = { {1, 32}, {32, 1}, @@ -602,6 +788,12 @@ std::vector processing_methods = { ProcessingMethod::CAST_ACT, }; +std::vector scaling_directions_2d_quantize = { + MXFP82DScalingDirection::RowwiseOnly, + MXFP82DScalingDirection::ColwiseOnly, + MXFP82DScalingDirection::Bidirectional, +}; + // Only GeLU activation tests are supported std::vector Activation_types = { ActivationType::Identity, @@ -621,6 +813,13 @@ class FusedCastMXFP8TestSuite : public ::testing::TestWithParam transformer_engine::DType, InputsFillCase>> {}; +class CastMXFP82DQuantizationTestSuite : public ::testing::TestWithParam + , + transformer_engine::DType, + transformer_engine::DType, + InputsFillCase>> {}; + TEST_P(FusedCastMXFP8TestSuite, TestFusedCastMXFP8) { #ifndef __HIP_PLATFORM_AMD__ // Skip tests for pre-Blackwell architectures @@ -703,6 +902,29 @@ TEST_P(FusedCastMXFP8TestSuite, TestFusedCastMXFP8) { } } +TEST_P(CastMXFP82DQuantizationTestSuite, TestCastMXFP82DQuantization) { + // Skip tests for pre-Blackwell architectures + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + + using namespace transformer_engine; + using namespace test; + + const auto scaling_direction = std::get<0>(GetParam()); + const auto matrix_size = std::get<1>(GetParam()); + const DType input_type = std::get<2>(GetParam()); + const DType output_type = std::get<3>(GetParam()); + const InputsFillCase fill_case = std::get<4>(GetParam()); + + TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY(input_type, InputType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8_ONLY(output_type, OutputType, + performTest_2d_quantize( + matrix_size, scaling_direction, fill_case); + ); + ); +} + std::string to_string(const ProcessingMethod method) { switch (method) { case ProcessingMethod::CAST_ONLY: return "CAST_ONLY"; @@ -726,6 +948,15 @@ std::string to_string(const ActivationType Act_type) { } } +std::string to_string(const MXFP82DScalingDirection scaling_direction) { + switch (scaling_direction) { + case MXFP82DScalingDirection::RowwiseOnly: return "RowwiseOnly"; + case MXFP82DScalingDirection::ColwiseOnly: return "ColwiseOnly"; + case MXFP82DScalingDirection::Bidirectional: return "Bidirectional"; + default: return ""; + } +} + std::string test_name_generator( const testing::TestParamInfo& info) { std::string name = to_string(std::get<0>(info.param)) + "X" + @@ -742,6 +973,19 @@ std::string test_name_generator( return name; } +std::string mxfp8_2d_quantization_test_name_generator( + const testing::TestParamInfo& info) { + std::string name = to_string(std::get<0>(info.param)); + const auto& shape = std::get<1>(info.param); + for ( const auto& s: shape) { + name += "X" + std::to_string(s); + } + name += "X" + test::typeName(std::get<2>(info.param)) + + "X" + test::typeName(std::get<3>(info.param)) + + "X" + test::caseName(std::get<4>(info.param)); + return name; +} + } // namespace // Test cases with only cast kernels @@ -758,6 +1002,18 @@ INSTANTIATE_TEST_SUITE_P( ::testing::ValuesIn(input_scenarios)), test_name_generator); +// Test cases for MXFP8 2D block scaling through the common C++ API. +INSTANTIATE_TEST_SUITE_P( + OperatorTest_CastMXFP8_2DQuantization, + CastMXFP82DQuantizationTestSuite, + ::testing::Combine( + ::testing::ValuesIn(scaling_directions_2d_quantize), + ::testing::ValuesIn(matrix_sizes_2d_quantize), + ::testing::Values(DType::kFloat32, DType::kBFloat16), + ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2), + ::testing::Values(InputsFillCase::uniform)), + mxfp8_2d_quantization_test_name_generator); + // Test cases with varying matrix shapes and block shapes INSTANTIATE_TEST_SUITE_P( OperatorTest_FusedCastMXFP8_Sizes, @@ -785,3 +1041,239 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2), ::testing::ValuesIn(input_scenarios)), test_name_generator); + +// ============================================================================ +// Swizzled-scales cast-only tests +// +// Validate the WITH_GEMM_SWIZZLED_SCALES=true code path added by the +// CastTraitsSwizzle port. The specialized kernel dispatches to +// CastTraitsSwizzle<..., kCacheColwise=true, kSwizzled=true> whenever the +// output tensor has set_with_gemm_swizzled_scales(true), producing scales +// directly in GEMM-swizzled layout. +// +// Reference construction: run the well-tested linear-scale path, then apply +// nvte_swizzle_scaling_factors (independently tested by SwizzleTestSuite) to +// transform the linear scales into the swizzled layout. Byte-compare the +// direct swizzled cast output against this reference for both scale tensors +// and the FP8 output data itself. +// +// Pass criteria per test: +// 1. FP8 rowwise data byte-identical between linear and swizzled paths. +// 2. FP8 colwise data byte-identical between linear and swizzled paths. +// 3. Swizzled rowwise scale bytes match linear->swizzle reference. +// 4. Swizzled colwise scale bytes match linear->swizzle reference. +// 5. Rowwise-only FP8 data and swizzled scales match the same reference. +// 6. No CUDA errors from any launch. +// ============================================================================ + +class SwizzledScalesFusedCastMXFP8TestSuite : public ::testing::TestWithParam< + std::tuple, + transformer_engine::DType, + transformer_engine::DType>> {}; + +TEST_P(SwizzledScalesFusedCastMXFP8TestSuite, TestSwizzledCastMXFP8) { + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + + using namespace transformer_engine; + using namespace test; + + const auto shape = std::get<0>(GetParam()); + const DType itype = std::get<1>(GetParam()); + const DType otype = std::get<2>(GetParam()); + + // BIDIMENSIONAL scaling needs a 2D+ input. + if (shape.size() < 2) { + GTEST_SKIP(); + } + + const size_t rows = first_dimension(shape); + const size_t cols = last_dimension(shape); + const size_t out_bytes = rows * cols; // fp8 = 1 byte/elem + + // Input filled with the same values for all casts. + Tensor input("input", shape, itype); + fillUniform(&input); + + // Target: swizzled-scale cast. Dispatcher routes to + // CastTraitsSwizzle<..., kCacheColwise=true, kSwizzled=true> on kernel #3. + Tensor output_swizzled("output_swizzled", shape, otype, true, true, NVTE_MXFP8_1D_SCALING); + output_swizzled.set_with_gemm_swizzled_scales(true); + nvte_quantize(input.data(), output_swizzled.data(), 0); + + // Rowwise-only activations take the pointer-based cast-only kernel. This + // directly exercises its WITH_GEMM_SWIZZLED_SCALES trait specialization + // for shapes whose column count is a multiple of 128. + Tensor output_rowwise_swizzled("output_rowwise_swizzled", shape, otype, + /*rowwise=*/true, /*colwise=*/false, + NVTE_MXFP8_1D_SCALING); + output_rowwise_swizzled.set_with_gemm_swizzled_scales(true); + nvte_quantize(input.data(), output_rowwise_swizzled.data(), 0); + + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << "swizzled-scale nvte_quantize failed"; + + // Reference construction: nvte_swizzle_scaling_factors accepts tensors with + // exactly one scale direction, so we build the rowwise and colwise references + // independently. Each is a single-direction linear cast followed by a + // single-direction swizzle transform, using the same input as the target. + // MXFP8 is a deterministic per-direction operation, so a rowwise-only cast + // produces the same rowwise fp8 bytes and scales as a rowwise+colwise cast. + + // Rowwise reference. + Tensor linear_row("linear_row", shape, otype, /*rowwise=*/true, /*colwise=*/false, + NVTE_MXFP8_1D_SCALING); + nvte_quantize(input.data(), linear_row.data(), 0); + Tensor ref_row_swz("ref_row_swz", shape, otype, /*rowwise=*/true, /*colwise=*/false, + NVTE_MXFP8_1D_SCALING); + ref_row_swz.set_with_gemm_swizzled_scales(true); + if (out_bytes > 0) { + cudaMemcpy(ref_row_swz.rowwise_dptr(), linear_row.rowwise_dptr(), + out_bytes, cudaMemcpyDeviceToDevice); + nvte_swizzle_scaling_factors(linear_row.data(), ref_row_swz.data(), 0); + } + + // Colwise reference. + Tensor linear_col("linear_col", shape, otype, /*rowwise=*/false, /*colwise=*/true, + NVTE_MXFP8_1D_SCALING); + nvte_quantize(input.data(), linear_col.data(), 0); + Tensor ref_col_swz("ref_col_swz", shape, otype, /*rowwise=*/false, /*colwise=*/true, + NVTE_MXFP8_1D_SCALING); + ref_col_swz.set_with_gemm_swizzled_scales(true); + if (out_bytes > 0) { + cudaMemcpy(ref_col_swz.columnwise_dptr(), linear_col.columnwise_dptr(), + out_bytes, cudaMemcpyDeviceToDevice); + nvte_swizzle_scaling_factors(linear_col.data(), ref_col_swz.data(), 0); + } + + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << "reference construction failed"; + + // ---- Comparisons ---- + + // (1) & (2): FP8 output data byte-identical to single-direction linear casts. + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8_ONLY(otype, OutputType, + { + const uint8_t *test_row = reinterpret_cast( + output_swizzled.rowwise_cpu_dptr()); + const uint8_t *ref_row = reinterpret_cast( + linear_row.rowwise_cpu_dptr()); + for (size_t i = 0; i < out_bytes; ++i) { + ASSERT_EQ(test_row[i], ref_row[i]) + << "rowwise fp8 data mismatch at index " << i + << " (swizzled=" << static_cast(test_row[i]) + << " linear=" << static_cast(ref_row[i]) << ")"; + } + const uint8_t *test_col = reinterpret_cast( + output_swizzled.columnwise_cpu_dptr()); + const uint8_t *ref_col = reinterpret_cast( + linear_col.columnwise_cpu_dptr()); + for (size_t i = 0; i < out_bytes; ++i) { + ASSERT_EQ(test_col[i], ref_col[i]) + << "colwise fp8 data mismatch at index " << i + << " (swizzled=" << static_cast(test_col[i]) + << " linear=" << static_cast(ref_col[i]) << ")"; + } + } + ); + + // (3): Swizzled rowwise scale bytes — directly-written vs linear-then-transform. + { + const size_t n = product(output_swizzled.rowwise_scale_inv_shape()); + const fp8e8m0 *test = output_swizzled.rowwise_cpu_scale_inv_ptr(); + const fp8e8m0 *ref = ref_row_swz.rowwise_cpu_scale_inv_ptr(); + for (size_t i = 0; i < n; ++i) { + ASSERT_EQ(test[i], ref[i]) + << "swizzled rowwise scale byte mismatch at offset " << i + << " (test=" << static_cast(test[i]) + << " ref=" << static_cast(ref[i]) << ")"; + } + } + + // (4): Swizzled colwise scale bytes — exercises the uint16-pair-packed flush + // path added by CACHE_COLWISE_SCALE_IN_SMEM + WITH_SWIZZLED_SCALES. + { + const size_t n = product(output_swizzled.columnwise_scale_inv_shape()); + const fp8e8m0 *test = output_swizzled.columnwise_cpu_scale_inv_ptr(); + const fp8e8m0 *ref = ref_col_swz.columnwise_cpu_scale_inv_ptr(); + for (size_t i = 0; i < n; ++i) { + ASSERT_EQ(test[i], ref[i]) + << "swizzled colwise scale byte mismatch at offset " << i + << " (test=" << static_cast(test[i]) + << " ref=" << static_cast(ref[i]) << ")"; + } + } + + // (5) & (6): The rowwise-only swizzled path matches the same independently + // constructed linear-then-swizzle reference. + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8_ONLY(otype, OutputType, + { + const uint8_t *test = reinterpret_cast( + output_rowwise_swizzled.rowwise_cpu_dptr()); + const uint8_t *ref = reinterpret_cast( + linear_row.rowwise_cpu_dptr()); + for (size_t i = 0; i < out_bytes; ++i) { + ASSERT_EQ(test[i], ref[i]) + << "rowwise-only fp8 data mismatch at index " << i; + } + } + ); + { + const size_t n = product(output_rowwise_swizzled.rowwise_scale_inv_shape()); + const fp8e8m0 *test = + output_rowwise_swizzled.rowwise_cpu_scale_inv_ptr(); + const fp8e8m0 *ref = ref_row_swz.rowwise_cpu_scale_inv_ptr(); + for (size_t i = 0; i < n; ++i) { + ASSERT_EQ(test[i], ref[i]) + << "rowwise-only swizzled scale mismatch at offset " << i; + } + } + +} + +std::string swizzled_test_name_generator( + const testing::TestParamInfo& info) { + std::string name; + const auto &shape = std::get<0>(info.param); + for (size_t i = 0; i < shape.size(); ++i) { + if (i > 0) name += "x"; + name += std::to_string(shape[i]); + } + name += "X" + test::typeName(std::get<1>(info.param)) + + "X" + test::typeName(std::get<2>(info.param)); + return name; +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest_FusedCastMXFP8_SwizzledCastOnly, + SwizzledScalesFusedCastMXFP8TestSuite, + ::testing::Values( + // 1. Aligned, small — sanity. + std::make_tuple(std::vector{128, 128}, DType::kBFloat16, DType::kFloat8E4M3), + // 2. Aligned, small — second dtype pair. + std::make_tuple(std::vector{128, 128}, DType::kFloat32, DType::kFloat8E5M2), + // 3. Aligned, medium. + std::make_tuple(std::vector{256, 384}, DType::kBFloat16, DType::kFloat8E4M3), + // 4. Odd number of scale rows (96/32 = 3) - exercises colwise flush + // odd-row scalar tail. + std::make_tuple(std::vector{96, 512}, DType::kBFloat16, DType::kFloat8E4M3), + // 5. cols/32 not a multiple of 4 - exercises rowwise flush scalar tail + // (need cols multiple of 8 for TMA-alignment on 16-bit input). + std::make_tuple(std::vector{256, 160}, DType::kBFloat16, DType::kFloat8E4M3), + // 6. Odd colwise scale rows + rowwise tail cols - combined tail paths. + std::make_tuple(std::vector{96, 160}, DType::kBFloat16, DType::kFloat8E4M3), + // 7. Small but tile-aligned. + std::make_tuple(std::vector{32, 64}, DType::kBFloat16, DType::kFloat8E4M3), + // 8. Minimum aligned size - single 32x32 tile. + std::make_tuple(std::vector{32, 32}, DType::kBFloat16, DType::kFloat8E4M3), + // 9. Multi-CTA at scale — stresses gmem writes across many CTAs. + std::make_tuple(std::vector{4096, 32768}, DType::kBFloat16, DType::kFloat8E4M3), + // 10. 4D input — matches the existing suite's rank coverage. + std::make_tuple(std::vector{16, 8, 4, 512}, DType::kBFloat16, DType::kFloat8E4M3), + // 11. Third input dtype. + std::make_tuple(std::vector{128, 128}, DType::kFloat16, DType::kFloat8E4M3), + // 12. Larger fp32. + std::make_tuple(std::vector{1024, 1024}, DType::kFloat32, DType::kFloat8E4M3) + ), + swizzled_test_name_generator); diff --git a/tests/cpp/operator/test_cast_mxfp8_grouped.cu b/tests/cpp/operator/test_cast_mxfp8_grouped.cu index 0144e698dc..b8dfdb7779 100644 --- a/tests/cpp/operator/test_cast_mxfp8_grouped.cu +++ b/tests/cpp/operator/test_cast_mxfp8_grouped.cu @@ -48,6 +48,7 @@ enum ShapeRepresentation { template void compute_ref(const ProcessingMethod processing_method, float (*OP)(const float), + const bool use_2d_quantization, const bool rowwise, const bool colwise, const InputType* input, @@ -122,13 +123,29 @@ void compute_ref(const ProcessingMethod processing_method, } } - if (rowwise) { + float block_amax_2d = 0.0f; + if (use_2d_quantization) { for (size_t i = i_min; i < i_max; ++i) { - float block_amax = 0.0f; - for (size_t j = j_min; j < j_max; ++j) { - const size_t cache_idx = (i - i_min) * tile_size_X + (j - j_min); - block_amax = std::max(block_amax, std::abs(cache_buffer[cache_idx])); + const size_t cache_idx = + (i - i_min) * tile_size_X + (j - j_min); + block_amax_2d = + std::max(block_amax_2d, std::abs(cache_buffer[cache_idx])); + } + } + } + + if (rowwise) { + for (size_t i = i_min; i < i_max; ++i) { + float block_amax = block_amax_2d; + + if (!use_2d_quantization) { + for (size_t j = j_min; j < j_max; ++j) { + const size_t cache_idx = + (i - i_min) * tile_size_X + (j - j_min); + block_amax = + std::max(block_amax, std::abs(cache_buffer[cache_idx])); + } } const fp8e8m0 biased_exponent = float_to_e8m0(block_amax * Quantized_Limits::max_reciprocal()); @@ -145,11 +162,15 @@ void compute_ref(const ProcessingMethod processing_method, } if (colwise) { for (size_t j = j_min; j < j_max; ++j) { - float block_amax = 0.0f; - - for (size_t i = i_min; i < i_max; ++i) { - const size_t cache_idx = (i - i_min) * tile_size_X + (j - j_min); - block_amax = std::max(block_amax, std::abs(cache_buffer[cache_idx])); + float block_amax = block_amax_2d; + + if (!use_2d_quantization) { + for (size_t i = i_min; i < i_max; ++i) { + const size_t cache_idx = + (i - i_min) * tile_size_X + (j - j_min); + block_amax = + std::max(block_amax, std::abs(cache_buffer[cache_idx])); + } } const fp8e8m0 biased_exponent = float_to_e8m0(block_amax * Quantized_Limits::max_reciprocal()); @@ -246,7 +267,8 @@ void performTest(const ProcessingMethod processing_method, const std::vector& last_dims_h, const std::vector& offsets_h, const bool rowwise, - const bool colwise) { + const bool colwise, + const bool use_2d_quantization = false) { using namespace test; DType itype = TypeInfo::dtype; @@ -503,7 +525,7 @@ void performTest(const ProcessingMethod processing_method, InputType* const ref_output_dbias_ptr = ref_output_dbias.data() + dbias_offset; compute_ref( - processing_method, OP, rowwise, colwise, in_ptr, grad_ptr, + processing_method, OP, use_2d_quantization, rowwise, colwise, in_ptr, grad_ptr, out_data_rowwise_ptr, out_data_colwise_ptr, out_scales_rowwise_ptr, out_scales_colwise_ptr, ref_output_dbias_ptr, M, K, @@ -512,6 +534,7 @@ void performTest(const ProcessingMethod processing_method, } QuantizationConfigWrapper quant_config; + quant_config.set_mxfp8_2d_quantization(use_2d_quantization); // GPU Tensor workspace; @@ -521,9 +544,9 @@ void performTest(const ProcessingMethod processing_method, break; } case ProcessingMethod::CAST_DBIAS: { - nvte_group_quantize_dbias(grad_group_tensor, out_group_tensor, output_dbias_tensor, workspace.data(), 0); + nvte_group_quantize_dbias(grad_group_tensor, out_group_tensor, output_dbias_tensor, workspace.data(), nullptr, 0); workspace = Tensor("workspace", workspace.rowwise_shape(), workspace.dtype()); - nvte_group_quantize_dbias(grad_group_tensor, out_group_tensor, output_dbias_tensor, workspace.data(), 0); + nvte_group_quantize_dbias(grad_group_tensor, out_group_tensor, output_dbias_tensor, workspace.data(), nullptr, 0); break; } case ProcessingMethod::CAST_DBIAS_DACT: { @@ -859,6 +882,35 @@ TEST_P(GroupedFusedCastMXFP8TestSuite, Test) { ); } +TEST(OperatorTest_GroupedFusedCastMXFP8, Test2DQuantization) { + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + + constexpr size_t num_tensors = 2; + const std::vector logical_shape = {256, 128}; + const std::vector first_dims = {128, 128}; + const std::vector last_dims = {128, 128}; + const std::vector offsets = {0, 128 * 128, 2 * 128 * 128}; + + for (const auto scaling_direction : scaling_directions) { + const bool rowwise = scaling_direction != ScalingDirection::COLWISE; + const bool colwise = scaling_direction != ScalingDirection::ROWWISE; + performTest( + ProcessingMethod::CAST_ONLY, + &identity, + ShapeRepresentation::SAME_BOTH_DIMS, + num_tensors, + logical_shape, + first_dims, + last_dims, + offsets, + rowwise, + colwise, + /*use_2d_quantization=*/true); + } +} + std::string to_string(const ProcessingMethod method) { switch (method) { case ProcessingMethod::CAST_ONLY: return "CAST_ONLY"; diff --git a/tests/cpp/operator/test_cast_nvfp4_transpose.cu b/tests/cpp/operator/test_cast_nvfp4_transpose.cu index 40c40d86f9..34c5783fcb 100644 --- a/tests/cpp/operator/test_cast_nvfp4_transpose.cu +++ b/tests/cpp/operator/test_cast_nvfp4_transpose.cu @@ -1417,6 +1417,138 @@ std::vector Activation_types = { ActivationType::Identity }; +// Element-level FP4 code differences between two compact NVFP4 buffers of the +// same logical (rows, cols) shape (cols must be even; FP4 is packed 2/byte). +inline size_t count_nvfp4_code_mismatches(const fp4e2m1* test_data, const fp4e2m1* ref_data, + int rows, int cols) { + size_t mismatches = 0; + for (int i = 0; i < rows; ++i) { + for (int j = 0; j < cols; j += 2) { + const int idx = i * cols + j; + const double2 t = cvt_fp4x2_to_double2(*reinterpret_cast(&test_data[idx / 2])); + const double2 r = cvt_fp4x2_to_double2(*reinterpret_cast(&ref_data[idx / 2])); + if (t.x != r.x) ++mismatches; + if (t.y != r.y) ++mismatches; + } + } + return mismatches; +} + +// Row-scaled transpose NVFP4 cast: emits rowwise (per-row amax) and columnwise +// (per-col amax) NVFP4 outputs in one pass. Cross-validated against the merged +// row-scaled 1D kernel (PR #2931, NVFP4ScalingMode::RowScaled1D): rowwise vs +// RowScaled1D(input), columnwise vs RowScaled1D(input^T). Amaxes and FP8 block +// scales must match bitwise; FP4 codes may differ only at rounding-midpoint +// ties (< 0.1%). bf16 input, rows/cols multiples of 128. +template +void performTestRowScaledTranspose(const std::vector& shape) { + using namespace test; + + const DType itype = TypeInfo::dtype; + const DType otype = DType::kFloat4E2M1; + const size_t rows = first_dimension(shape); + const size_t cols = last_dimension(shape); + + Tensor input("input", shape, itype); + fillCase(&input, InputsFillCase::uniform); + + // System under test: row-scaled NVFP4 with both directions requested in one + // call. The transpose kernel is selected by the generic nvte_quantize_v2 + // dispatch when a row-scaled tensor (set_row_scaled_nvfp4) also allocates a + // columnwise output -- no dedicated config flag. + Tensor output("output", shape, otype, /*rowwise=*/true, /*columnwise=*/true, + NVTE_NVFP4_1D_SCALING); + // Marking the tensor row-scaled with a columnwise output allocated selects the + // transpose kernel and sizes the per-row (rowwise) and per-col (columnwise) + // amax vectors (default is a single scalar amax). + output.set_row_scaled_nvfp4(true); + QuantizationConfigWrapper quant_config; + quant_config.set_stochastic_rounding(false); + nvte_quantize_v2(input.data(), output.data(), quant_config, 0); + + // Reference (rowwise direction): trusted row-scaled 1D kernel on the input. + Tensor ref_row("ref_row", shape, otype, /*rowwise=*/true, /*columnwise=*/false, + NVTE_NVFP4_1D_SCALING); + ref_row.set_row_scaled_nvfp4(true); + QuantizationConfigWrapper ref_config; + ref_config.set_stochastic_rounding(false); + nvte_quantize_v2(input.data(), ref_row.data(), ref_config, 0); + + // Reference (columnwise direction): trusted row-scaled 1D kernel on input^T. + input.to_cpu(); + std::vector input_t_host = + create_transpose(input.rowwise_cpu_dptr(), rows, cols); + const std::vector shape_t = {cols, rows}; + Tensor input_t("input_t", shape_t, itype); + std::copy(input_t_host.begin(), input_t_host.end(), input_t.rowwise_cpu_dptr()); + input_t.from_cpu(); + Tensor ref_col("ref_col", shape_t, otype, /*rowwise=*/true, /*columnwise=*/false, + NVTE_NVFP4_1D_SCALING); + ref_col.set_row_scaled_nvfp4(true); + nvte_quantize_v2(input_t.data(), ref_col.data(), ref_config, 0); + + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); + + output.to_cpu(); + ref_row.to_cpu(); + ref_col.to_cpu(); + + // FP4 codes must match the trusted kernel except at rounding-midpoint ties. + // Cap the tolerated disagreement well below what a real bug would produce. + constexpr double kMaxCodeMismatchRate = 5e-3; // 0.5% (observed < 0.1%) + const size_t row_mismatches = count_nvfp4_code_mismatches( + output.rowwise_cpu_dptr(), ref_row.rowwise_cpu_dptr(), + static_cast(rows), static_cast(cols)); + EXPECT_LT(static_cast(row_mismatches) / static_cast(rows * cols), + kMaxCodeMismatchRate) + << "rowwise FP4 disagreement " << row_mismatches << "/" << (rows * cols); + const size_t col_mismatches = count_nvfp4_code_mismatches( + output.columnwise_cpu_dptr(), ref_col.rowwise_cpu_dptr(), + static_cast(cols), static_cast(rows)); + EXPECT_LT(static_cast(col_mismatches) / static_cast(cols * rows), + kMaxCodeMismatchRate) + << "columnwise FP4 disagreement " << col_mismatches << "/" << (cols * rows); + + // FP8 e4m3 block scale factors must match (compact layout on both sides). + const std::array sd = get_scale_tensor_dims(rows, cols, 1, 16); + const std::array sd_t = get_scale_tensor_dims(cols, rows, 1, 16); + size_t scale_mismatches = 0; + compare_scaling_factors("rowwise_scales", + output.rowwise_cpu_scale_inv_ptr(), + ref_row.rowwise_cpu_scale_inv_ptr(), + sd[0], sd[1], sd[3], scale_mismatches); + compare_scaling_factors("columnwise_scales", + output.columnwise_cpu_scale_inv_ptr(), + ref_col.rowwise_cpu_scale_inv_ptr(), + sd_t[0], sd_t[1], sd_t[3], scale_mismatches); + ASSERT_EQ(scale_mismatches, 0u); + + // Per-row / per-col amaxes must match the reference amaxes exactly. + ASSERT_EQ(output.rowwise_amax_size(), rows); + const float* row_amax = output.cpu_rowwise_amax_ptr(); + const float* ref_row_amax = ref_row.cpu_rowwise_amax_ptr(); + for (size_t i = 0; i < rows; ++i) { + ASSERT_EQ(row_amax[i], ref_row_amax[i]) << "rowwise amax mismatch at row " << i; + } + const float* col_amax = output.cpu_columnwise_amax_ptr(); + const float* ref_col_amax = ref_col.cpu_rowwise_amax_ptr(); + for (size_t j = 0; j < cols; ++j) { + ASSERT_EQ(col_amax[j], ref_col_amax[j]) << "columnwise amax mismatch at col " << j; + } +} + +// Row-scaled transpose requires 128-aligned rows and cols (bf16 only). +std::vector> row_scaled_transpose_dims = { + {128, 128}, + {256, 256}, + {128, 256}, + {256, 512}, + {512, 512}, + {384, 1024}, + {2048, 256}, +}; + } // namespace class FusedCastTransposeNVFP4TestSuite : public ::testing::TestWithParam @@ -1599,3 +1731,25 @@ INSTANTIATE_TEST_SUITE_P( } return name; }); + +class CastNVFP4RowScaledTransposeTestSuite : public ::testing::TestWithParam> {}; + +TEST_P(CastNVFP4RowScaledTransposeTestSuite, MatchesRowScaledReference) { + // The row-scaled transpose NVFP4 cast kernel requires Blackwell. + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + performTestRowScaledTranspose(GetParam()); +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, + CastNVFP4RowScaledTransposeTestSuite, + ::testing::ValuesIn(row_scaled_transpose_dims), + [](const testing::TestParamInfo& info) { + std::string name; + for (const auto& s : info.param) { + name += "X" + std::to_string(s); + } + return name; + }); diff --git a/tests/cpp/test_common.cu b/tests/cpp/test_common.cu index ea217f749d..cff249a532 100644 --- a/tests/cpp/test_common.cu +++ b/tests/cpp/test_common.cu @@ -515,13 +515,17 @@ void Tensor::set_row_scaled_nvfp4(bool row_scaled_nvfp4) { // Update amax tensor if (row_scaled_nvfp4) { - // Row-scaled NVFP4 has amax matching number of rows NVTE_CHECK(rowwise_, "Row-scaled NVFP4 requires row-wise data."); - NVTE_CHECK(!columnwise_, "Row-scaled NVFP4 does not support column-wise data."); auto shape = tensor_.shape(); const size_t rows = product(shape, 0, shape.ndim - 1); + const size_t cols = shape.data[shape.ndim - 1]; amax_rowwise_.emplace(rows, DType::kFloat32); tensor_.set_amax(amax_rowwise_->gpu_buffer(), DType::kFloat32, std::vector{rows}); + if (columnwise_) { + amax_columnwise_.emplace(cols, DType::kFloat32); + tensor_.set_columnwise_amax(amax_columnwise_->gpu_buffer(), DType::kFloat32, + std::vector{cols}); + } } else { // Tensor-scaled NVFP4 has single amax if (rowwise_) { diff --git a/tests/cpp_distributed/test_ep.cu b/tests/cpp_distributed/test_ep.cu index 7dbbcdce9d..be6dbed3ab 100644 --- a/tests/cpp_distributed/test_ep.cu +++ b/tests/cpp_distributed/test_ep.cu @@ -97,6 +97,26 @@ static std::vector expected_recv_values_sorted( return vals; } +// Sorted multiset of per-token identity bytes each recv_rank should receive, +// derived from the deterministic routing map (same id_of stamp as the test). +static std::vector expected_recv_ids_sorted( + int recv_rank, int num_processes, int num_tokens, int top_k, + int num_experts, int num_local_experts) { + int base = recv_rank * num_local_experts; + std::vector ids; + for (int src = 0; src < num_processes; ++src) { + auto idx = routing_balanced(src, num_tokens, top_k, num_experts, num_local_experts); + for (int t = 0; t < num_tokens; ++t) + for (int k = 0; k < top_k; ++k) { + int64_t e = idx[t * top_k + k]; + if (e >= base && e < base + num_local_experts) + ids.push_back(static_cast((src * num_tokens + t + 1) & 0xFF)); + } + } + std::sort(ids.begin(), ids.end()); + return ids; +} + // 2^-5 relative tolerance for BF16 (matches mantissa precision with margin), // plus a small atol floor for near-zero expected values. static constexpr float kBf16Rtol = 1.0f / 32.0f; @@ -368,6 +388,123 @@ TYPED_TEST(EPDispatchTest, PrepareAndDispatch) { NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); } +// ============================================================================= +// EPDispatchScaledTest: MXFP8 dispatch routes each token's e8m0 scale row in +// lockstep with its e4m3 data row. +// ============================================================================= + +// Each source token stamps a nonzero identity byte across its whole data row +// and its whole scale row. Dispatch is a pure permutation, so every filled recv +// slot must carry the same id in both buffers; a zero id means scales were not +// routed. +TEST_F(EpOpTestBase, MXFP8DispatchScales) { + if (g_sm_major < 9) GTEST_SKIP() << "EP requires SM_90+"; + // MXFP8 dispatch is supported on all HT EM modes (local_permute, local_dup, + // nvlink_dup); the mode is selected by the NCCL_EP_HT_EM_* env at group init. + // MXFP8 e8m0 scale bytes/token = hidden/32; NCCL EP HT requires 16B alignment. + if (hidden_dim_ % 512 != 0) + GTEST_SKIP() << "MXFP8 dispatch needs hidden_dim % 512 == 0 (16B scale alignment)"; + const int scale_cols = hidden_dim_ / 32; + using Shape = std::vector; + + EPBuffers buf; + buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, + ep_size_, max_tokens_per_rank_); + EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + + auto h_idx = routing_balanced(g_process_id, num_tokens_, top_k_, + num_experts_, num_local_experts_); + std::vector h_w(num_tokens_ * top_k_, 1.0f / top_k_); + NVTE_CHECK_CUDA(cudaMemcpy(buf.topk_idx.get(), h_idx.data(), + h_idx.size() * sizeof(int64_t), cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA(cudaMemcpy(buf.topk_weights.get(), h_w.data(), + h_w.size() * sizeof(float), cudaMemcpyHostToDevice)); + + auto id_of = [&](int tok) { + return static_cast((g_process_id * num_tokens_ + tok + 1) & 0xFF); + }; + std::vector h_data(num_tokens_ * hidden_dim_); + std::vector h_scale(num_tokens_ * scale_cols); + for (int tok = 0; tok < num_tokens_; ++tok) { + const uint8_t id = id_of(tok); + std::fill_n(&h_data[tok * hidden_dim_], hidden_dim_, id); + std::fill_n(&h_scale[tok * scale_cols], scale_cols, id); + } + + DevBuf d_data(num_tokens_ * hidden_dim_), d_scale(num_tokens_ * scale_cols); + DevBuf d_recv_data(buf.recv_capacity * hidden_dim_); + DevBuf d_recv_scale(buf.recv_capacity * scale_cols); + NVTE_CHECK_CUDA(cudaMemcpy(d_data.get(), h_data.data(), h_data.size(), cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA(cudaMemcpy(d_scale.get(), h_scale.data(), h_scale.size(), cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA(cudaMemset(d_recv_data.get(), 0, d_recv_data.bytes())); + NVTE_CHECK_CUDA(cudaMemset(d_recv_scale.get(), 0, d_recv_scale.bytes())); + + TensorWrapper tokens(NVTE_MXFP8_1D_SCALING); + tokens.set_rowwise_data(d_data.get(), DType::kFloat8E4M3, + Shape{(size_t)num_tokens_, (size_t)hidden_dim_}); + tokens.set_rowwise_scale_inv(d_scale.get(), DType::kFloat8E8M0, + Shape{(size_t)num_tokens_, (size_t)scale_cols}); + TensorWrapper recv_tokens(NVTE_MXFP8_1D_SCALING); + recv_tokens.set_rowwise_data(d_recv_data.get(), DType::kFloat8E4M3, + Shape{buf.recv_capacity, (size_t)hidden_dim_}); + recv_tokens.set_rowwise_scale_inv(d_recv_scale.get(), DType::kFloat8E8M0, + Shape{buf.recv_capacity, (size_t)scale_cols}); + + cudaStream_t stream; + NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); + ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), + t.recv_tokens_per_expert.data(), nullptr, &t.layer_cfg_, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(t.handle_mem.data(), t.topk_idx.data(), + tokens.data(), NVTECommWindow{}, t.topk_weights.data(), + NVTECommWindow{}, recv_tokens.data(), NVTECommWindow{}, + t.recv_topk_weights.data(), NVTECommWindow{}, stream)); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + + std::vector counts(num_local_experts_); + NVTE_CHECK_CUDA(cudaMemcpy(counts.data(), buf.recv_tokens_per_expert.get(), + num_local_experts_ * sizeof(int32_t), cudaMemcpyDeviceToHost)); + auto exp_counts = expected_recv_tokens_per_expert(g_process_id, g_num_processes, num_tokens_, + top_k_, num_experts_, num_local_experts_); + int total_recv = 0; + for (int e = 0; e < num_local_experts_; ++e) { + EXPECT_EQ(counts[e], exp_counts[e]) << "local expert " << e; + total_recv += exp_counts[e]; + } + ASSERT_LE(total_recv, static_cast(buf.recv_capacity)); + + std::vector r_data(buf.recv_capacity * hidden_dim_); + std::vector r_scale(buf.recv_capacity * scale_cols); + NVTE_CHECK_CUDA(cudaMemcpy(r_data.data(), d_recv_data.get(), r_data.size(), cudaMemcpyDeviceToHost)); + NVTE_CHECK_CUDA(cudaMemcpy(r_scale.data(), d_recv_scale.get(), r_scale.size(), cudaMemcpyDeviceToHost)); + + std::vector got_data_ids, got_scale_ids; + got_data_ids.reserve(total_recv); + got_scale_ids.reserve(total_recv); + size_t slot = 0; + for (int e = 0; e < num_local_experts_; ++e) { + for (int i = 0; i < counts[e]; ++i, ++slot) { + got_data_ids.push_back(r_data[slot * hidden_dim_]); + got_scale_ids.push_back(r_scale[slot * scale_cols]); + } + } + std::sort(got_data_ids.begin(), got_data_ids.end()); + std::sort(got_scale_ids.begin(), got_scale_ids.end()); + + auto exp_ids = expected_recv_ids_sorted(g_process_id, g_num_processes, num_tokens_, + top_k_, num_experts_, num_local_experts_); + ASSERT_EQ(got_data_ids.size(), exp_ids.size()); + ASSERT_EQ(got_scale_ids.size(), exp_ids.size()); + for (size_t i = 0; i < exp_ids.size(); ++i) { + EXPECT_EQ(got_data_ids[i], exp_ids[i]) << "recv data id mismatch at sorted index " << i; + EXPECT_EQ(got_scale_ids[i], exp_ids[i]) << "recv scale id mismatch at sorted index " << i; + } + + if (g_process_id == 0) + printf(" MXFP8DispatchScales: passed (recv=%d, data + scales match routing map)\n", total_recv); + + NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); +} + // ============================================================================= // EPCombineTest: round-trip identity expert -> result == top_k * tokens. // ============================================================================= @@ -675,7 +812,7 @@ class EPPipelineTest : public EpOpTestBase, public ::testing::WithParamInterface TEST_P(EPPipelineTest, FullForwardBackward) { const DType dtype = GetParam(); // NCCL EP backend currently asserts ncclBfloat16 in ncclEpDispatch - // (contrib/nccl_ep/nccl_ep.cc); skip FP16/FP32 until the backend supports them. + // (nccl_ep/nccl_ep.cc); skip FP16/FP32 until the backend supports them. if (dtype != DType::kBFloat16) { GTEST_SKIP() << test::typeName(dtype) << " not yet supported by NCCL EP backend"; } @@ -750,8 +887,9 @@ class EPZeroCopyTest : public EpOpTestBase { }; TYPED_TEST_SUITE(EPZeroCopyTest, EPBf16Only); -// Identity round-trip with symm-mem on dispatch i/o + combine input. Bit-exact -// vs HBM reference (same routing, same input). +// Identity round-trip with symm-mem on dispatch i/o + combine input. The combined +// result is bit-exact vs the HBM reference; the intermediate recv buffer is not, +// since zero-copy and HBM dispatch use different per-expert layouts. TYPED_TEST(EPZeroCopyTest, IdentityAllSymm) { using Tok = TypeParam; EP_PULL_FIXTURE(); @@ -776,10 +914,7 @@ TYPED_TEST(EPZeroCopyTest, IdentityAllSymm) { ref_t.result.data(), stream)); NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); - std::vector ref_recv(ref_buf.recv_capacity * hidden_dim_); std::vector ref_result(num_tokens_ * hidden_dim_); - NVTE_CHECK_CUDA(cudaMemcpy(ref_recv.data(), ref_buf.recv_tokens.get(), - ref_recv.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); NVTE_CHECK_CUDA(cudaMemcpy(ref_result.data(), ref_buf.result.get(), ref_result.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); @@ -818,24 +953,17 @@ TYPED_TEST(EPZeroCopyTest, IdentityAllSymm) { symm_window(sym_recv), sym_t.result.data(), stream)); NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); - std::vector sym_recv_host(sym_buf.recv_capacity * hidden_dim_); std::vector sym_result(num_tokens_ * hidden_dim_); - NVTE_CHECK_CUDA(cudaMemcpy(sym_recv_host.data(), sym_recv.ptr, - sym_recv_host.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); NVTE_CHECK_CUDA(cudaMemcpy(sym_result.data(), sym_buf.result.get(), sym_result.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); - // Compare per filled recv slot (HBM ref vs symm) and full result. - int total_recv = this->template read_total_recv(sym_buf); - for (int i = 0; i < total_recv * hidden_dim_; ++i) - ASSERT_EQ(tok_to_float(sym_recv_host[i]), tok_to_float(ref_recv[i])) - << "recv mismatch at " << i; + // Combined result is the cross-mode invariant (see note above). for (size_t i = 0; i < sym_result.size(); ++i) ASSERT_EQ(tok_to_float(sym_result[i]), tok_to_float(ref_result[i])) << "result mismatch at " << i; if (g_process_id == 0) - printf(" IdentityAllSymm: passed (recv_slots=%d, bit-exact vs HBM)\n", total_recv); + printf(" IdentityAllSymm: passed (result bit-exact vs HBM)\n"); NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); } diff --git a/tests/cpp_distributed/test_ep_common.h b/tests/cpp_distributed/test_ep_common.h index 7cf6017090..37f324eabb 100644 --- a/tests/cpp_distributed/test_ep_common.h +++ b/tests/cpp_distributed/test_ep_common.h @@ -7,7 +7,7 @@ /* * Shared TE EP test infrastructure. Include once per TU; ep_bootstrap() in * each test binary's main() populates process-level globals. - * Defaults: 4 experts/rank, hidden_dim=256, max_tokens_per_rank=64. + * Defaults: 4 experts/rank, hidden_dim=512, max_tokens_per_rank=64. */ #pragma once @@ -48,7 +48,7 @@ static int g_num_processes = -1; static int g_sm_major = -1; // set by ep_bootstrap; -1 until then static int g_ep_size = -1; static int g_num_experts = -1; -static int g_hidden_dim = 256; +static int g_hidden_dim = 512; static int g_max_tokens_per_rank = 64; static NVTEDType g_max_token_dtype = kNVTEFloat32; // staging-buffer sizing static bool g_ep_initialized = false; diff --git a/tests/jax/distributed_test_base.py b/tests/jax/distributed_test_base.py index 6d963f5c7b..c65ffbf62b 100644 --- a/tests/jax/distributed_test_base.py +++ b/tests/jax/distributed_test_base.py @@ -17,6 +17,17 @@ def generate_configs(): configs = [] + if is_devices_enough(8): + configs.append( + pytest.param( + 8, + (2, 2, 2), + ("dp", "fsdp", "tpsp"), + MeshResource(dp_resource="dp", fsdp_resource="fsdp", tpsp_resource="tpsp"), + id="n8_dp2_fsdp2_tp2", + ) + ) + if is_devices_enough(4): configs.append( pytest.param( @@ -79,6 +90,7 @@ def generate_collectives_count(allreduce, allgather, other): def assert_equal_collectives(target_hlo, coll_count_ref): target_splitted_hlo = target_hlo.splitlines() start_symb = "-start" + sync_symb = '"is_sync":true' def count_bytes(hlo_text): bytes_count = 0 @@ -114,15 +126,62 @@ def get_bytes_per_txt(t): return bytes_count + def get_called_collective_type(line): + """Identify a collective hidden inside an async/fusion wrapper.""" + match = re.search(r"\bcalls=(%[-.\w]+)", line) + if not match: + return None + + computation_name = match.group(1) + computation = re.search( + rf"(?ms)^[ \t]*{re.escape(computation_name)}(?=[ \t(])" rf".*?^[ \t]*}}[ \t]*$", + target_hlo, + ) + if not computation: + return None + + computation_text = computation.group(0) + has_all_reduce = COLL_AR_KEY in computation_text + has_all_gather = COLL_AG_KEY in computation_text + + if has_all_reduce and not has_all_gather: + return COLL_AR_KEY + if has_all_gather and not has_all_reduce: + return COLL_AG_KEY + + return None + def count_collectives(splitted_hlo): result = generate_collectives_count(0, 0, 0) for line in splitted_hlo: txt = line.split() - if len(txt) > 0 and start_symb in txt[0]: - if COLL_AR_KEY in txt[0]: + + # strip optional HLO syntax prefix + if txt and txt[0] == "ROOT": + txt = txt[1:] + + # Asynchronous collectives are represented by *-start and *-done + # instructions, so count only *-start. Synchronous collectives are + # represented by a single instruction without either suffix. + is_async_start = txt and start_symb in txt[0] + is_sync_collective = "collective_backend_config" in line and sync_symb in line + + if is_async_start or is_sync_collective: + # Some direct *-start instructions are tagged with + # `"is_sync":true`. Classify the explicit async form first so + # it is not mistaken for an unsuffixed synchronous collective. + if is_async_start: + called_collective = get_called_collective_type(line) + is_all_reduce = COLL_AR_KEY in txt[0] or called_collective == COLL_AR_KEY + is_all_gather = COLL_AG_KEY in txt[0] or called_collective == COLL_AG_KEY + else: + is_all_reduce = re.search(r"\ball-reduce\s*\(", line) + is_all_gather = re.search(r"\ball-gather\s*\(", line) + + if is_all_reduce: result[COLL_AR_KEY] += count_bytes(txt) - elif COLL_AG_KEY in txt[0]: + elif is_all_gather: result[COLL_AG_KEY] += count_bytes(txt) else: result[COLL_OTHER_KEY] += count_bytes(txt) diff --git a/tests/jax/multi_process_launch_ep.sh b/tests/jax/multi_process_launch_ep.sh index ff89f712eb..8547d77f2b 100755 --- a/tests/jax/multi_process_launch_ep.sh +++ b/tests/jax/multi_process_launch_ep.sh @@ -17,8 +17,8 @@ SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" export PYTHONPATH="${TE_REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" -# Editable installs don't embed rpath; libtransformer_engine.so needs -# libnccl_ep.so.0 from the TE editable location at dlopen time. +# Editable installs don't embed rpath; the TE JAX extension needs +# libtransformer_engine.so from the TE editable location at dlopen time. TE_LIB_PATH=$(pip3 show transformer-engine 2>/dev/null \ | grep -E "Location:|Editable project location:" \ | tail -n 1 | awk '{print $NF}') diff --git a/tests/jax/run_te_ep_moe.sh b/tests/jax/run_te_ep_moe.sh index 32d5f21956..33352d3b3a 100755 --- a/tests/jax/run_te_ep_moe.sh +++ b/tests/jax/run_te_ep_moe.sh @@ -14,7 +14,7 @@ set -euo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" TE_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" -TEST_FILE="$TE_ROOT/tests/jax/test_te_ep_moe.py" +TEST_FILE="${TEST_FILE:-$TE_ROOT/tests/jax/test_te_ep_moe.py}" PYTEST_INI="$TE_ROOT/tests/jax/pytest.ini" NUM_GPUS="${NUM_GPUS:-$(nvidia-smi -L | wc -l)}" diff --git a/tests/jax/test_distributed_dense.py b/tests/jax/test_distributed_dense.py index 7aed9dafcc..8b118d48be 100644 --- a/tests/jax/test_distributed_dense.py +++ b/tests/jax/test_distributed_dense.py @@ -50,6 +50,9 @@ def _get_sharding_for_gemm(mesh, mesh_resource, partition_layout="rowwise"): dp_axis = mesh_resource.dp_resource tp_axis = mesh_resource.tpsp_resource + if mesh_resource.fsdp_resource is not None: + dp_axis = (mesh_resource.dp_resource, mesh_resource.fsdp_resource) + if partition_layout == "colwise": x_spec = PartitionSpec(dp_axis, None, None) weight_spec = PartitionSpec(None, tp_axis) diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index e1b965a54e..f1609ad09b 100644 --- a/tests/jax/test_distributed_fused_attn.py +++ b/tests/jax/test_distributed_fused_attn.py @@ -26,6 +26,7 @@ _has_cudnn_frontend_python, ) from utils import pytest_parametrize_wrapper +from transformer_engine_jax import get_cudnn_version, get_device_compute_capability from transformer_engine.jax.attention import ( is_fused_attn_kernel_available, AttnBiasType, @@ -349,6 +350,65 @@ def test_softcap_score_mod_with_aux_params_backward( pytest.param([4, 256, 16, 64], id="4-256xCPx2-16-64"), ] +DISTRIBUTED_CONTEXT_SELF_ATTN_D256_DATA_SHAPES = { + "L0": [], + "L1": [[2, 128, 16, 256]], + "L2": [], +} + +# Keep these as explicit tuples instead of independent layout/mask/window as: +# BSHD CP uses CAUSAL_MASK, THD CP uses PADDING_CAUSAL_MASK, SWA is +# only valid for THD, and stripe_size behavior is different for +# BSHD vs THD in these tests. Splitting the axes would mostly add +# invalid BSHD+SWA and THD+CAUSAL combinations that fail or skip. +DISTRIBUTED_CONTEXT_SELF_ATTN_D256_LAYOUTS_MASKS_WINDOWS = [ + # BSHD with different layouts, but same causal mask and no sliding window + pytest.param( + QKVLayout.BSHD_BS2HD, + AttnMaskType.CAUSAL_MASK, + (-1, -1), + id="BSHD_KVPACKED-CAUSAL-NO_SWA", + ), + pytest.param( + QKVLayout.BSHD_BSHD_BSHD, + AttnMaskType.CAUSAL_MASK, + (-1, -1), + id="BSHD_SEPARATE-CAUSAL-NO_SWA", + ), + # THD with different sliding window sizes, but same packed layout and padding causal mask + pytest.param( + QKVLayout.THD_T2HD, + AttnMaskType.PADDING_CAUSAL_MASK, + (-1, -1), + id="THD_KVPACKED-PADDING_CAUSAL-NO_SWA", + ), + pytest.param( + QKVLayout.THD_T2HD, + AttnMaskType.PADDING_CAUSAL_MASK, + (20, 0), + id="THD_KVPACKED-PADDING_CAUSAL-SWA", + ), + # THD with different sliding window sizes, but same separate layout and padding causal mask + pytest.param( + QKVLayout.THD_THD_THD, + AttnMaskType.PADDING_CAUSAL_MASK, + (-1, -1), + id="THD_SEPARATE-PADDING_CAUSAL-NO_SWA", + ), + pytest.param( + QKVLayout.THD_THD_THD, + AttnMaskType.PADDING_CAUSAL_MASK, + (20, 0), + id="THD_SEPARATE-PADDING_CAUSAL-SWA", + ), +] + +DISTRIBUTED_CONTEXT_SELF_ATTN_MAX_LOGIT_CP_MODES = [ + pytest.param(CPStrategy.ALL_GATHER, False, id="AG"), + pytest.param(CPStrategy.RING, False, id="RING-NO_SCAN"), + pytest.param(CPStrategy.RING, True, id="RING-SCAN"), +] + class TestDistributedContextParallelSelfAttn: # TODO(KshitijLakhani): parametrize num_segments_per_seq for all CP tests @@ -369,6 +429,8 @@ def impl_test_context_parallel_attn( window_size=None, stripe_size=None, num_segments_per_seq=None, + return_max_logit=False, + check_forward_output=True, ): if qkv_layout.is_thd(): if not load_balanced and ( @@ -470,9 +532,89 @@ def check_has_backend_for_mask(mask_type): if attn_mask_type == AttnMaskType.CAUSAL_MASK and mesh_shape[1] != 1: #CP pytest.skip(f"Skipping CAUSAL_MASK and CP={mesh_shape[1]} for AOTriton") - runner.test_backward() + if return_max_logit: + runner.test_forward( + return_max_logit=True, + check_output=check_forward_output, + ) + else: + runner.test_backward() del os.environ["NVTE_FUSED_RING_ATTENTION_USE_SCAN"] + @pytest_parametrize_wrapper( + "device_count,mesh_shape,mesh_axes,mesh_resource", + generate_context_parallel_configs_for_attn(), + ) + @pytest.mark.parametrize("data_shape", DISTRIBUTED_CONTEXT_SELF_ATTN_DATA_SHAPES[:1]) + @pytest.mark.parametrize("kv_groups", [1, 8]) + @pytest.mark.parametrize("dtype", [pytest.param(jnp.bfloat16, id="BF16")]) + @pytest.mark.parametrize( + "qkv_layout, attn_mask_type", + DISTRIBUTED_CONTEXT_SELF_ATTN_LAYOUTS_MASKS, + ) + @pytest.mark.parametrize( + "cp_strategy, use_scan_ring", + DISTRIBUTED_CONTEXT_SELF_ATTN_MAX_LOGIT_CP_MODES, + ) + @pytest.mark.parametrize( + "window_size", + [ + pytest.param((-1, -1), id="NO_SWA"), + pytest.param((20, 0), id="SWA"), + ], + ) + def test_context_parallel_return_max_logit( + self, + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + data_shape, + kv_groups, + dtype, + qkv_layout, + attn_mask_type, + cp_strategy, + window_size, + use_scan_ring, + ): + """Check CP fused attention returns global per-head max_logit.""" + is_thd = qkv_layout.is_thd() + supports_swa = is_thd and ( + cp_strategy == CPStrategy.ALL_GATHER + or (cp_strategy == CPStrategy.RING and not use_scan_ring) + ) + if window_size != (-1, -1) and not supports_swa: + pytest.skip("CP SWA requires THD All-Gather or unrolled THD Ring.") + # TODO: Evaluate cuDNN Max mismatches observed for striped multi-segment THD Ring GQA. + if is_thd and cp_strategy == CPStrategy.RING and kv_groups > 1: + pytest.skip("THD Ring GQA Max mismatches require further evaluation.") + + stripe_size = 64 if is_thd and cp_strategy == CPStrategy.ALL_GATHER else None + if is_thd and cp_strategy == CPStrategy.RING: + stripe_size = 1 + num_segments_per_seq = 5 if is_thd else None + check_forward_output = not (is_thd and cp_strategy == CPStrategy.RING) + self.impl_test_context_parallel_attn( + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + data_shape, + kv_groups, + attn_mask_type, + dtype, + qkv_layout, + True, + cp_strategy, + use_scan_ring=use_scan_ring, + window_size=window_size, + stripe_size=stripe_size, + num_segments_per_seq=num_segments_per_seq, + return_max_logit=True, + check_forward_output=check_forward_output, + ) + @pytest_parametrize_wrapper( "device_count,mesh_shape,mesh_axes,mesh_resource", generate_context_parallel_configs_for_attn(), @@ -496,7 +638,7 @@ def check_has_backend_for_mask(mask_type): "window_size", [ pytest.param((-1, -1), id="window_size(-1, -1)"), - pytest.param((5, 0), id="window_size(8, 0)"), + pytest.param((5, 0), id="window_size(5, 0)"), ], ) @pytest.mark.parametrize( @@ -651,6 +793,119 @@ def test_context_parallel_ring_attn( stripe_size=stripe_size, ) + # CP ring and all-gather tests for D=256 + # TODO(KshitijLakhani): Replace this with common-provided fused-attn disable reasons once + # they can be surfaced to framework tests. + @staticmethod + def skip_if_d256_cp_unsupported(qkv_layout): + compute_capability = get_device_compute_capability(0) + if not 100 <= compute_capability < 110: + pytest.skip("D=256 CP fused attention is only enabled on Blackwell server GPUs.") + + required_cudnn_version = 92500 if qkv_layout.is_thd() else 92300 + required_cudnn_version_label = "9.25" if qkv_layout.is_thd() else "9.23" + if get_cudnn_version() < required_cudnn_version: + pytest.skip( + f"D=256 CP fused attention with {qkv_layout} requires cuDNN" + f" {required_cudnn_version_label} or newer." + ) + + @pytest_parametrize_wrapper( + "device_count,mesh_shape,mesh_axes,mesh_resource", + generate_context_parallel_configs_for_attn(), + ) + @pytest_parametrize_wrapper( + "data_shape", + DISTRIBUTED_CONTEXT_SELF_ATTN_D256_DATA_SHAPES, + ) + @pytest.mark.parametrize( + "dtype", + [pytest.param(jnp.float16, id="FP16"), pytest.param(jnp.bfloat16, id="BF16")], + ) + @pytest.mark.parametrize( + "qkv_layout, attn_mask_type, window_size", + DISTRIBUTED_CONTEXT_SELF_ATTN_D256_LAYOUTS_MASKS_WINDOWS, + ) + def test_context_parallel_ring_attn_d256( + self, + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + data_shape, + dtype, + qkv_layout, + attn_mask_type, + window_size, + ): + """D=256 CP ring coverage.""" + self.skip_if_d256_cp_unsupported(qkv_layout) + + self.impl_test_context_parallel_attn( + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + data_shape, + 1, + attn_mask_type, + dtype, + qkv_layout, + True, + CPStrategy.RING, + use_scan_ring=False, + window_size=window_size, + stripe_size=1 if qkv_layout.is_thd() else None, + ) + + @pytest_parametrize_wrapper( + "device_count,mesh_shape,mesh_axes,mesh_resource", + generate_context_parallel_configs_for_attn(), + ) + @pytest_parametrize_wrapper( + "data_shape", + DISTRIBUTED_CONTEXT_SELF_ATTN_D256_DATA_SHAPES, + ) + @pytest.mark.parametrize( + "dtype", + [pytest.param(jnp.float16, id="FP16"), pytest.param(jnp.bfloat16, id="BF16")], + ) + @pytest.mark.parametrize( + "qkv_layout, attn_mask_type, window_size", + DISTRIBUTED_CONTEXT_SELF_ATTN_D256_LAYOUTS_MASKS_WINDOWS, + ) + def test_context_parallel_allgather_attn_d256( + self, + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + data_shape, + dtype, + qkv_layout, + attn_mask_type, + window_size, + ): + """D=256 CP all-gather coverage.""" + self.skip_if_d256_cp_unsupported(qkv_layout) + + self.impl_test_context_parallel_attn( + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + data_shape, + 1, + attn_mask_type, + dtype, + qkv_layout, + True, + CPStrategy.ALL_GATHER, + window_size=window_size, + stripe_size=128 if qkv_layout.is_thd() else None, + num_segments_per_seq=5 if qkv_layout.is_thd() else None, + ) + REORDER_CAUSAL_LOAD_BALANCING_DATA_SHAPES = { "L0": [[]], @@ -681,7 +936,7 @@ def test(self, cp_size, shape, qkv_format, reorder_strategy, stripe_size): seq_dim = 0 if reorder_strategy == ReorderStrategy.Striped: - seq_lens = shape[seq_dim] + seq_lens = tensor.shape[seq_dim] if seq_lens < (cp_size * stripe_size): pytest.skip(f"{seq_lens=} must be larger than {cp_size*stripe_size=}") diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 13b861969e..cafba71f1d 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -67,7 +67,7 @@ def init(): yield -@partial(jax.jit, static_argnums=(6, 7, 8, 9, 11, 12)) +@partial(jax.jit, static_argnums=(6, 7, 8, 9, 11, 12, 13)) def general_dot_product_attention( query: ArrayLike, key: ArrayLike, @@ -82,25 +82,22 @@ def general_dot_product_attention( dropout_rng: ArrayLike, dtype: DTypeLike, score_mod_reference: Optional[Callable[[Array], Array]] = None, + is_max_logit_enabled: bool = False, ) -> Array: """ Similar to flax.linen.dot_product_attention but with GQA support """ query, key, value, bias = promote_dtype(query, key, value, bias, dtype=dtype) dtype = query.dtype - b, s_q, h_q, d = query.shape _, s_kv, h_kv, _ = key.shape assert (h_q % h_kv == 0) and (h_q >= h_kv) num_groups = h_q // h_kv grouped_query = jnp.reshape(query, (b, s_q, h_kv, num_groups, d)) - # logits with shape (b, h_kv, num_groups, s_q, s_kv) logits = scale_factor * jnp.einsum("...qhgd,...khd->...hgqk", grouped_query, key) if bias is not None: - # reshape logits without groups logits = logits.reshape((b, h_kv * num_groups, s_q, s_kv)) - # apply post-scale bias logits = logits + bias # [ROCm] Detect query rows where ALL bias values are -inf (fully masked out). # These rows would produce NaN in softmax; zero logits to prevent NaN since @@ -120,6 +117,8 @@ def general_dot_product_attention( if score_mod_reference is not None: # Kernel tests use NO_MASK; fused_attn rejects mask+score_mod before this reference path. logits = score_mod_reference(logits.astype(jnp.float32)) + if is_max_logit_enabled: + return jnp.max(logits.reshape((b, h_q, s_q, s_kv)), axis=(0, 2, 3)).astype(dtype) match softmax_type: case AttnSoftmaxType.VANILLA_SOFTMAX: @@ -287,7 +286,17 @@ def _split_valid_and_invalid(primitive, reference, pad): return primitive_valid, primitive_invalid, reference_valid, reference_invalid -def jax_dpa(query, key, value, bias, softmax_offset, mask, dropout_rng, **kwargs): +def jax_dpa( + query, + key, + value, + bias, + softmax_offset, + mask, + dropout_rng, + is_max_logit_enabled=False, + **kwargs, +): """ JAX native dot product attention implementation """ @@ -327,6 +336,7 @@ def jax_dpa(query, key, value, bias, softmax_offset, mask, dropout_rng, **kwargs dropout_rng=dropout_rng, dtype=jnp.float32, score_mod_reference=score_mod_reference, + is_max_logit_enabled=is_max_logit_enabled, ) return output.astype(query.dtype) @@ -358,9 +368,13 @@ def customcall_fused_dpa( qkv_args = (query, key, value) case _: raise ValueError(f"Unsupported {qkv_layout=}") - return fused_attn( + result = fused_attn( qkv_args, bias, sequence_descriptor, dropout_rng, softmax_offset=softmax_offset, **kwargs - ).astype(query.dtype) + ) + if isinstance(result, tuple): + output, max_logit = result + return output.astype(query.dtype), max_logit + return result.astype(query.dtype) @pytest.mark.skipif( @@ -492,19 +506,19 @@ def _get_max_segments_per_sequence(self): return 1 def _check_configs(self): - # TODO(rewang): probably adds this in is_fused_attn_available + # TODO(KshitijLakhani): probably add/move this to is_fused_attn_available if self.qkv_layout.is_thd() and not self.attn_mask_type.is_padding(): pytest.skip("THD format requires padding masks.") if self.attn_mask_type.is_bottom_right(): if self.max_seqlen_q > self.max_seqlen_kv: pytest.skip( - f"BRCM requires cross attn type pattern, i.e.max_seqlen_kv >= max_seqlen_q" + "BRCM requires cross attn type pattern, i.e.max_seqlen_kv >= max_seqlen_q" ) if self.attn_bias_type is not AttnBiasType.NO_BIAS: - pytest.skip(f"cuDNN does not support pre or post scale bias for BRCM") + pytest.skip("cuDNN does not support pre or post scale bias for BRCM") if self.dropout_prob != 0.0: - pytest.skip(f"cuDNN does not support non-zero dropoouts for BRCM") + pytest.skip("cuDNN does not support non-zero dropouts for BRCM") if self.qkv_layout.is_qkvpacked(): if self.max_seqlen_q != self.max_seqlen_kv: @@ -516,11 +530,62 @@ def _check_configs(self): pytest.skip( "seqlen_q > seqlen_kv is not supported with sliding window attention in cuDNN" ) + compute_capability = get_device_compute_capability(0) + cudnn_version = get_cudnn_version() + # D=256 bprop on SM10x uses the deterministic algorithm path only. BSHD support + # starts with cuDNN FE 1.24 / BE 9.23; THD execution-plan support starts with + # cuDNN FE 1.26 / BE 9.25. The kernel rejects dBias, dropout, and ALiBi, supports vanilla + # softmax only, and allows SWA together with a causal mask only. + is_sm10x = 100 <= compute_capability < 110 + if self.is_training and is_sm10x and (self.head_dim_qk == 256 or self.head_dim_v == 256): + if self.head_dim_qk != 256 or self.head_dim_v != 256: + pytest.skip( + "D=256 BWD on Blackwell only supports d_qk == d_v == 256;" + f" got d_qk={self.head_dim_qk}, d_v={self.head_dim_v}." + ) + required_cudnn_version = 92500 if self.qkv_layout.is_thd() else 92300 + required_cudnn_version_label = "9.25" if self.qkv_layout.is_thd() else "9.23" + if cudnn_version < required_cudnn_version: + pytest.skip( + f"D=256 BWD on Blackwell with {self.qkv_layout} requires cuDNN" + f" {required_cudnn_version_label} or newer; got cuDNN {cudnn_version}." + ) + # TODO(KshitijLakhani): cuDNN FE can model bias input separately from dBias, + # but TE does not yet plumb whether dBias is requested into the common backend selector. + # Until that distinction is available, the D=256 SM10x gate requires no bias. + unsupported = None + if self.attn_bias_type == AttnBiasType.PRE_SCALE_BIAS: + unsupported = "pre-scale bias" + elif self.attn_bias_type != AttnBiasType.NO_BIAS: + unsupported = ( + "post-scale bias in TE's D=256 backend gate; bias-input-only" + " support needs TE to distinguish between bias input and dBias" + ) + elif self.dropout_prob != 0.0: + unsupported = "dropout" + elif self.softmax_type != AttnSoftmaxType.VANILLA_SOFTMAX: + unsupported = "non-vanilla softmax" + if unsupported is not None: + pytest.skip( + "D=256 BWD on Blackwell uses the deterministic SM100 D=256 SDPA BWD" + f" kernel which does not support {unsupported}." + ) + if self.window_size is not None and self.window_size != (-1, -1): + if not self.attn_mask_type.is_causal(): + pytest.skip( + "D=256 BWD on Blackwell uses the SM100 D=256 SDPA BWD kernel" + " which requires window_size=(-1, -1) for non-causal masks." + ) + if self.window_size[1] not in (-1, 0): + pytest.skip( + "D=256 BWD on Blackwell only supports right window -1 or 0" + " for causal masks." + ) - if not is_hip_extension() and get_device_compute_capability(0) >= 100 and self.is_training: + if not is_hip_extension() and compute_capability >= 100 and self.is_training: if FusedAttnHelper.is_non_deterministic_allowed() and ( (self.dropout_prob != 0.0 and self.attn_bias_type != AttnBiasType.NO_BIAS) - or get_cudnn_version() < 90700 + or cudnn_version < 90700 ): pytest.skip( "For sm100+, non-deterministic bprop (cuDNN 9.7+) does not support bias with" @@ -529,7 +594,7 @@ def _check_configs(self): if not FusedAttnHelper.is_non_deterministic_allowed() and ( self.dropout_prob != 0.0 or self.attn_bias_type != AttnBiasType.NO_BIAS - or get_cudnn_version() < 91801 + or cudnn_version < 91801 ): pytest.skip( "For sm100+, deterministic bprop (cuDNN 9.18.1+) does not support bias or" @@ -949,7 +1014,7 @@ def to_dp_shardings(x): self.seq_length_offset_pspec = PartitionSpec(self.mesh_resource.dp_resource, None) self.seq_length_offset_sharding = NamedSharding(self.mesh, self.seq_length_offset_pspec) - def test_forward(self): + def test_forward(self, return_max_logit=False, check_output=True): """ Test forward with JITted primitive and unJITted reference """ @@ -997,6 +1062,7 @@ def test_forward(self): "score_mod_bprop": self.score_mod_bprop, "score_mod_tensors": self.score_mod_tensors, "score_mod_bprop_tensors": self.score_mod_bprop_tensors, + "return_max_logit": return_max_logit, } reference_kwargs = {**kwargs, "score_mod_reference": self.score_mod_reference} @@ -1016,31 +1082,37 @@ def test_forward(self): with jax.set_mesh(self.mesh), autocast(mesh_resource=self.mesh_resource): primitive_out = customcall_fused_dpa_jit(*customcall_args) + if return_max_logit: + primitive_out, primitive_max_logit = primitive_out primitive_out = self.cp_inverse_reorder_fn(primitive_out) - reference_out = jax_dpa(*args, **reference_kwargs) + if return_max_logit: + reference_max_logit = jax_dpa(*args, is_max_logit_enabled=True, **reference_kwargs) - if self.is_training and self.dropout_prob > 0.0: - return + if check_output and not (self.is_training and self.dropout_prob > 0.0): + reference_out = jax_dpa(*args, **reference_kwargs) - primitive_valid, primitive_invalid, reference_valid, reference_invalid = ( - _split_valid_and_invalid(primitive_out, reference_out, self.pad_q) - ) + primitive_valid, primitive_invalid, reference_valid, _ = _split_valid_and_invalid( + primitive_out, reference_out, self.pad_q + ) - assert_allclose( - primitive_invalid, - jnp.zeros_like(primitive_invalid), - rtol=self.rtol, - atol=self.atol, - dtype=self.dtype, - ) - assert_allclose( - primitive_valid, - reference_valid, - rtol=self.rtol, - atol=self.atol, - dtype=self.dtype, - ) + assert_allclose( + primitive_invalid, + jnp.zeros_like(primitive_invalid), + rtol=self.rtol, + atol=self.atol, + dtype=self.dtype, + ) + assert_allclose( + primitive_valid, + reference_valid, + rtol=self.rtol, + atol=self.atol, + dtype=self.dtype, + ) + + if return_max_logit: + assert_allclose(primitive_max_logit, reference_max_logit, dtype=self.dtype) if self.coll_count_ref is not None: with jax.set_mesh(self.mesh), autocast(mesh_resource=self.mesh_resource): @@ -1049,7 +1121,7 @@ def test_forward(self): ) assert_equal_collectives(target_hlo, self.coll_count_ref) - def test_backward(self): + def test_backward(self, return_max_logit=False): """ Test value_and_grad with JIT, which includes both forward and backward. @@ -1077,6 +1149,8 @@ def grad_func( if self.attn_mask_type.is_causal(): gradient_multiplier /= 10 output = func(q, k, v, bias, softmax_offset, sequence_descriptor, dropout_rng, **kwargs) + if isinstance(output, tuple): + output, _ = output if cp_reverse_out: output = self.cp_inverse_reorder_fn(output) # Keep only valid result for the gradient @@ -1142,6 +1216,7 @@ def grad_func( "score_mod_bprop": self.score_mod_bprop, "score_mod_tensors": self.score_mod_tensors, "score_mod_bprop_tensors": self.score_mod_bprop_tensors, + "return_max_logit": return_max_logit, } reference_kwargs = {**kwargs, "score_mod_reference": self.score_mod_reference} @@ -1309,6 +1384,123 @@ def check_dqkv(primitive, reference, pad, idx): assert_equal_collectives(target_hlo, self.coll_count_ref) +FUSED_ATTN_MAX_LOGIT_QKV_LAYOUTS = [ + pytest.param( + QKVLayout.BSHD_BSHD_BSHD, + 8, + 8, + id="BSHD_SEPARATE", + ), + pytest.param( + QKVLayout.BS3HD, + 8, + 8, + id="BS3HD", + ), + pytest.param( + QKVLayout.BSHD_BS2HD, + 8, + 4, + id="BSHD_KV_PACKED-GQA", + ), + pytest.param( + QKVLayout.T3HD, + 8, + 8, + id="THD_QKV_PACKED", + ), + pytest.param( + QKVLayout.THD_THD_THD, + 8, + 8, + id="THD_SEPARATE", + ), +] + + +class TestFusedAttnMaxLogit: + """Targeted non-CP max_logit coverage.""" + + @staticmethod + @pytest.mark.parametrize( + "qkv_layout, num_heads_q, num_heads_kv", + FUSED_ATTN_MAX_LOGIT_QKV_LAYOUTS, + ) + @pytest.mark.parametrize( + "attn_bias_type, bias_shape", + [ + pytest.param(AttnBiasType.NO_BIAS, None, id="NO_BIAS"), + pytest.param( + AttnBiasType.POST_SCALE_BIAS, + BiasShape._1HSS, + id="POST_SCALE_BIAS-1HSS", + ), + ], + ) + @pytest.mark.parametrize( + "attn_mask_type", + [ + pytest.param(AttnMaskType.NO_MASK, id="NO_MASK"), + pytest.param(AttnMaskType.PADDING_MASK, id="PADDING_MASK"), + pytest.param(AttnMaskType.CAUSAL_MASK, id="CAUSAL_MASK"), + pytest.param(AttnMaskType.PADDING_CAUSAL_MASK, id="PADDING_CAUSAL_MASK"), + ], + ) + def test_forward( + qkv_layout, + num_heads_q, + num_heads_kv, + attn_bias_type, + bias_shape, + attn_mask_type, + ): + """Check non-CP JAX fused attention can expose framework-compatible max_logit.""" + runner = FusedAttnRunner( + batch_size=2, + max_seqlen_q=128, + max_seqlen_kv=128, + num_heads_q=num_heads_q, + num_heads_kv=num_heads_kv, + head_dim_qk=64, + head_dim_v=64, + attn_bias_type=attn_bias_type, + attn_mask_type=attn_mask_type, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + dropout_prob=0.0, + dtype=jnp.bfloat16, + is_training=True, + qkv_layout=qkv_layout, + bias_shape=bias_shape, + window_size=None, + seq_desc_format=SeqDescFormat.Seqlens, + ) + runner.test_forward(return_max_logit=True) + + @staticmethod + def test_backward(): + """Ensure aux-return cotangents do not break the fused attention backward path.""" + runner = FusedAttnRunner( + batch_size=2, + max_seqlen_q=128, + max_seqlen_kv=128, + num_heads_q=8, + num_heads_kv=8, + head_dim_qk=64, + head_dim_v=64, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.PADDING_CAUSAL_MASK, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + dropout_prob=0.0, + dtype=jnp.bfloat16, + is_training=True, + qkv_layout=QKVLayout.BSHD_BSHD_BSHD, + bias_shape=None, + window_size=None, + seq_desc_format=SeqDescFormat.Seqlens, + ) + runner.test_backward(return_max_logit=True) + + def _get_swa_window_size_for_test(s_kv: int, attn_mask_type: AttnMaskType) -> Tuple[int, int]: """Pick a sliding-window size for SWA tests, gated on cuDNN version. @@ -1790,6 +1982,32 @@ def test_backward( QKVLayout.THD_THD_THD, id="2-1024-2048-12-6-128-64-BF16-CROSS-GQA-RAGGED_SEPARATE", ), + # D=256 deterministic backward on the SM100 dedicated SDPA bprop kernel. + # BSHD requires cuDNN FE 1.24 / BE 9.23+; THD requires cuDNN FE 1.26 / BE 9.25+. + pytest.param( + 4, + 128, + 128, + 16, + 16, + 256, + 256, + jnp.float16, + QKVLayout.BSHD_BS2HD, + id="4-128-128-16-16-256-256-FP16-SELF-KV_PACKED", + ), + pytest.param( + 4, + 128, + 128, + 16, + 16, + 256, + 256, + jnp.float16, + QKVLayout.THD_T2HD, + id="4-128-128-16-16-256-256-FP16-SELF-RAGGED_KV_PACKED", + ), ], ) @pytest.mark.parametrize( diff --git a/tests/jax/test_fused_attn_score_mod.py b/tests/jax/test_fused_attn_score_mod.py index b1f165f491..7d1137b21f 100644 --- a/tests/jax/test_fused_attn_score_mod.py +++ b/tests/jax/test_fused_attn_score_mod.py @@ -426,6 +426,7 @@ def fake_fused_attn( score_mod_bprop=None, score_mod_tensors=None, score_mod_bprop_tensors=None, + return_max_logit=False, ): captured.update( qkv=qkv, diff --git a/tests/jax/test_multi_process_ep.py b/tests/jax/test_multi_process_ep.py index 0b8bb25f3f..47af0b0c39 100644 --- a/tests/jax/test_multi_process_ep.py +++ b/tests/jax/test_multi_process_ep.py @@ -14,11 +14,15 @@ - ``ep_dispatch`` custom_vjp: exact per-(t, k) ``grad_topk_weights`` under skewed upstream gradients (no k-axis averaging). - HLO reshard guard: compile-only, no XLA collectives outside the EP FFI. + - Drop-on-overflow: ``ep_finalize`` + re-bootstrap with a small recv capacity + drops the overflow instead of trapping; ``total_recv_tokens`` reports the + pre-drop demand. Launch via tests/jax/multi_process_launch_ep.sh (one process per GPU). """ import os +import re import sys import unittest @@ -28,14 +32,23 @@ import numpy as np from jax.sharding import Mesh, NamedSharding, PartitionSpec +from utils import is_devices_enough from transformer_engine.jax.sharding import MeshResource, global_shard_guard -from transformer_engine.jax.ep import EpLayerConfig, ep_bootstrap, ep_dispatch, ep_combine +from transformer_engine.jax.ep import ( + EpLayerConfig, + ep_bootstrap, + ep_finalize, + ep_dispatch, + ep_combine, + _ep_domain_for_rank, +) from transformer_engine.jax.cpp_extensions.ep import ( ep_prepare, ep_dispatch_fwd, ep_combine_fwd, get_ep_config, ) +from transformer_engine.jax.version_utils import is_collective_stream_supported # ── Test config ───────────────────────────────────────────────────────────── @@ -230,8 +243,8 @@ def test_two_handle_mems_no_aliasing(self): @jax.jit def run(idx): - _tc_a, ha = ep_prepare(ka, idx) - _tc_b, hb = ep_prepare(kb, idx) + _tc_a, _trt_a, ha = ep_prepare(ka, idx) + _tc_b, _trt_b, hb = ep_prepare(kb, idx) return ha, hb hm_a, hm_b = run(idx_s) @@ -258,7 +271,9 @@ def test_two_layer_dispatch_no_handle_aliasing(self): w = jax.lax.with_sharding_constraint(topk_w, NamedSharding(self.mesh, dp_spec)) def one_layer(hk, idx, toks, w_): - recv_t, recv_w, hm, tc = ep_dispatch(hk, idx, toks, w_, self.recv_capacity_per_rank) + recv_t, recv_w, hm, tc, _trt = ep_dispatch( + hk, idx, toks, w_, self.recv_capacity_per_rank + ) recv_t = jax.lax.with_sharding_constraint( recv_t, NamedSharding(self.mesh, ep_spec_3d) ) @@ -296,23 +311,33 @@ def run(idx, ta_, tb_, w_): ) def test_primitive_prepare(self): - """ep_prepare returns token_counts and handle_mem of the expected shapes.""" + """ep_prepare returns token_counts, total_recv_tokens and handle_mem. + + total_recv_tokens is the per-rank pre-drop recv-slot total; with no + overflow it equals the padded per-expert count sum that dispatch fills. + """ T_global, topk_idx, _tokens, _w = self._make_identity_inputs() del T_global dp_spec = PartitionSpec(("dp", "ep"), None) + align = max(int(self.hk.dispatch_output_per_expert_alignment), 1) with self.mesh, global_shard_guard(self.mr): idx_s = jax.lax.with_sharding_constraint(topk_idx, NamedSharding(self.mesh, dp_spec)) @jax.jit def run(idx): - tc, hm = ep_prepare(self.hk, idx) - return tc, hm + tc, trt, hm = ep_prepare(self.hk, idx) + return tc, trt, hm - tc, hm = run(idx_s) - tc.block_until_ready() + tc, trt, hm = run(idx_s) + trt.block_until_ready() self.assertEqual(tc.shape, (self.dp * self.ep, NUM_LOCAL_EXPERTS)) self.assertEqual(hm.shape[0], self.dp * self.ep) self.assertGreater(hm.shape[1], 0) + self.assertEqual(trt.shape, (self.dp * self.ep, 1)) + self.assertEqual(trt.dtype, jnp.int32) + padded = ((np.asarray(tc).astype(np.int64) + align - 1) // align) * align + expected = padded.sum(axis=-1, keepdims=True) + np.testing.assert_array_equal(np.asarray(trt).astype(np.int64), expected) def _run_identity_round_trip(self, nonuniform): T_global, topk_idx, tokens, topk_w = self._make_identity_inputs(nonuniform=nonuniform) @@ -327,7 +352,7 @@ def _run_identity_round_trip(self, nonuniform): @jax.jit def run(idx, toks, w): - _tc, hm = ep_prepare(self.hk, idx) + _tc, _trt, hm = ep_prepare(self.hk, idx) recv_t, recv_w = ep_dispatch_fwd( self.hk, hm, idx, toks, w, self.recv_capacity_per_rank ) @@ -393,7 +418,7 @@ def loss_fn(toks): toks = jax.lax.with_sharding_constraint(toks, NamedSharding(self.mesh, dp_spec)) idx = jax.lax.with_sharding_constraint(topk_idx, NamedSharding(self.mesh, dp_spec)) w = jax.lax.with_sharding_constraint(topk_w, NamedSharding(self.mesh, dp_spec)) - recv_t, recv_w, hm, tc = ep_dispatch( + recv_t, recv_w, hm, tc, _trt = ep_dispatch( self.hk, idx, toks, w, self.recv_capacity_per_rank ) recv_t = jax.lax.with_sharding_constraint( @@ -444,7 +469,7 @@ def test_dispatch_combine_3d_input_output(self): @jax.jit def run(idx, toks, w): - recv_t, recv_w, hm, _tc = ep_dispatch( + recv_t, recv_w, hm, _tc, _trt = ep_dispatch( self.hk, idx, toks, w, self.recv_capacity_per_rank ) recv_t = jax.lax.with_sharding_constraint(recv_t, NamedSharding(self.mesh, ep_t)) @@ -495,7 +520,7 @@ def loss_fn(toks): toks = jax.lax.with_sharding_constraint(toks, NamedSharding(self.mesh, dp_spec)) idx = jax.lax.with_sharding_constraint(topk_idx, NamedSharding(self.mesh, dp_spec)) w = jax.lax.with_sharding_constraint(topk_w, NamedSharding(self.mesh, dp_spec)) - recv_tokens, _recv_w, _hm, tc = ep_dispatch( + recv_tokens, _recv_w, _hm, tc, _trt = ep_dispatch( self.hk, idx, toks, w, self.recv_capacity_per_rank ) recv_tokens = jax.lax.with_sharding_constraint( @@ -549,7 +574,7 @@ def loss_fn(eo): toks = jax.lax.with_sharding_constraint(tokens, NamedSharding(self.mesh, dp_spec)) idx = jax.lax.with_sharding_constraint(topk_idx, NamedSharding(self.mesh, dp_spec)) w = jax.lax.with_sharding_constraint(topk_w, NamedSharding(self.mesh, dp_spec)) - _recv_tokens, recv_w, hm, tc = ep_dispatch( + _recv_tokens, recv_w, hm, tc, _trt = ep_dispatch( self.hk, idx, toks, w, self.recv_capacity_per_rank ) recv_w = jax.lax.with_sharding_constraint( @@ -595,7 +620,7 @@ def loss_fn(idx_in, tok_in, w_in): idx_in = jax.lax.with_sharding_constraint(idx_in, NamedSharding(self.mesh, dp_spec)) tok_in = jax.lax.with_sharding_constraint(tok_in, NamedSharding(self.mesh, dp_spec)) w_in = jax.lax.with_sharding_constraint(w_in, NamedSharding(self.mesh, dp_spec)) - _recv_t, recv_w, _h, _tc = ep_dispatch( + _recv_t, recv_w, _h, _tc, _trt = ep_dispatch( self.hk, idx_in, tok_in, w_in, self.recv_capacity_per_rank ) # Per-slot index scale ⇒ each slot's contribution differs. @@ -636,7 +661,7 @@ def run(idx, toks, w): idx = jax.lax.with_sharding_constraint(idx, NamedSharding(self.mesh, dp_spec)) toks = jax.lax.with_sharding_constraint(toks, NamedSharding(self.mesh, dp_spec)) w = jax.lax.with_sharding_constraint(w, NamedSharding(self.mesh, dp_spec)) - recv_t, recv_w, hm, tc = ep_dispatch( + recv_t, recv_w, hm, tc, _trt = ep_dispatch( self.hk, idx, toks, w, self.recv_capacity_per_rank ) recv_t = jax.lax.with_sharding_constraint( @@ -660,6 +685,87 @@ def run(idx, toks, w): expected = (("dp", "ep"),) if self.dp > 1 else ("ep",) self.assertEqual(tuple(compiled.output_shardings.spec), expected) + @unittest.skipUnless( + is_collective_stream_supported(), + "JAX/XLA lacks the gpu_stream:collective annotation (openxla/xla#39604)", + ) + def test_z_dispatch_combine_on_collective_stream(self): + """Every EP FFI custom call must run on the collective stream. compute_on + puts the annotation on the async wrapper XLA generates, so assert each EP + call is reachable from a wrapper that carries it.""" + T_dp, tokens, topk_idx, topk_w = self._make_random_inputs() + dp_spec = PartitionSpec(("dp", "ep"), None) + ep_spec_3d = PartitionSpec(("dp", "ep"), None, None) + ep_spec_2d = PartitionSpec(("dp", "ep"), None) + + with self.mesh, global_shard_guard(self.mr): + + @jax.jit + def run(idx, toks, w): + idx = jax.lax.with_sharding_constraint(idx, NamedSharding(self.mesh, dp_spec)) + toks = jax.lax.with_sharding_constraint(toks, NamedSharding(self.mesh, dp_spec)) + w = jax.lax.with_sharding_constraint(w, NamedSharding(self.mesh, dp_spec)) + recv_t, recv_w, hm, tc, _trt = ep_dispatch( + self.hk, idx, toks, w, self.recv_capacity_per_rank + ) + recv_t = jax.lax.with_sharding_constraint( + recv_t, NamedSharding(self.mesh, ep_spec_3d) + ) + recv_w = jax.lax.with_sharding_constraint( + recv_w, NamedSharding(self.mesh, ep_spec_2d) + ) + weighted = self._preweight_expert_out(recv_t, recv_w) + out = ep_combine(self.hk, hm, tc, weighted, T_dp, out_sharding=(("dp", "ep"), None)) + return jax.lax.with_sharding_constraint(out, NamedSharding(self.mesh, dp_spec)) + + hlo = run.lower(topk_idx, tokens, topk_w).compile().as_text() + + # Parse the HLO into computations and follow the call graph: a call + # carrying the collective-stream annotation places its callee (and every + # nested callee) on the collective stream. + comps = {} + cur = None + for line in hlo.splitlines(): + stripped = line.strip() + header = re.match(r"(?:ENTRY\s+)?(%[\w.\-]+)\s*\(", stripped) + if header and stripped.endswith("{"): + cur = header.group(1) + comps[cur] = [] + elif stripped == "}": + cur = None + elif cur is not None: + comps[cur].append(line) + + callees = lambda l: re.findall(r"(?:calls|to_apply)=(%[\w.\-]+)", l) + edges = {c: {x for l in ls for x in callees(l)} for c, ls in comps.items()} + + collective = set() + for ls in comps.values(): + for l in ls: + if '_xla_stream_annotation="collective"' in l.replace(" ", ""): + collective.update(callees(l)) + stack = list(collective) + while stack: + for callee in edges.get(stack.pop(), ()): + if callee not in collective: + collective.add(callee) + stack.append(callee) + + ep_calls = [ + (c, l) for c, ls in comps.items() for l in ls if 'custom_call_target="te_ep_' in l + ] + self.assertTrue(ep_calls, f"no te_ep_* custom calls in compiled HLO:\n{hlo}") + missing = [ + l.strip()[:200] + for c, l in ep_calls + if c not in collective + and '_xla_stream_annotation="collective"' not in l.replace(" ", "") + ] + self.assertFalse( + missing, + "te_ep_* custom calls not on the collective stream:\n" + "\n".join(missing), + ) + def test_z_no_unexpected_reshard_in_hlo_bwd(self): """Compiled bwd HLO must not insert XLA collectives outside the EP FFI.""" T_dp, tokens, topk_idx, topk_w = self._make_random_inputs() @@ -682,7 +788,9 @@ def fwd(eo, toks, idx, w): toks = jax.lax.with_sharding_constraint(toks, NamedSharding(self.mesh, dp_spec)) idx = jax.lax.with_sharding_constraint(idx, NamedSharding(self.mesh, dp_spec)) w = jax.lax.with_sharding_constraint(w, NamedSharding(self.mesh, dp_spec)) - _rt, rw, hm, tc = ep_dispatch(self.hk, idx, toks, w, self.recv_capacity_per_rank) + _rt, rw, hm, tc, _trt = ep_dispatch( + self.hk, idx, toks, w, self.recv_capacity_per_rank + ) rw = jax.lax.with_sharding_constraint(rw, NamedSharding(self.mesh, ep_spec_2d)) weighted = self._preweight_expert_out(eo, rw) combined = ep_combine( @@ -712,6 +820,147 @@ def bwd_only(eo, toks, idx, w, g): self.assertEqual(hlo.count(op), 0, f"unexpected XLA {op} in bwd HLO:\n{hlo}") +# ── Drop-on-overflow ───────────────────────────────────────────────────────── + + +class TestEPOverflowDrop(unittest.TestCase): + """Re-bootstraps with drop_on_overflow=True and a small recv capacity. + + ``ep_finalize`` tears down the default TestEP communicator so this class can + re-bootstrap with its own config in the same process. Every token routes its + top-1 slot to expert 0, so the rank owning it demands more than the recv + capacity; the dispatch drops the excess instead of trapping, while + total_recv_tokens still reports the pre-drop demand. + """ + + ALIGN = 16 + # Each EP group routes all OVF_TOKENS_PER_DP_SHARD top-1 slots to expert 0, + # whose padded count then exceeds the recv capacity. HT mode requires the + # recv capacity to be >= the per-rank dispatch count (tokens // ep), so the + # capacity sits between that floor and the concentrated expert-0 demand. + OVF_TOKENS_PER_DP_SHARD = 48 + OVF_RECV_CAPACITY = NUM_LOCAL_EXPERTS * ALIGN # 32 slots per rank + + @classmethod + def setUpClass(cls): + sm = _local_device_sm() + if sm is not None and sm < 90: + raise unittest.SkipTest(f"NCCL EP requires SM>=90 (got SM{sm})") + cls.num_procs = jax.process_count() + cls.rank = jax.process_index() + cls.dp, cls.ep = _factor_dp_ep(cls.num_procs) + cls.num_experts = NUM_LOCAL_EXPERTS * cls.ep + cls.recv_capacity_per_rank = cls.OVF_RECV_CAPACITY + # True per-rank dispatch count; HT mode needs recv_capacity >= this. + cls.max_tokens_per_rank = cls.OVF_TOKENS_PER_DP_SHARD // cls.ep + cls.mesh = _build_mesh(cls.dp, cls.ep) + cls.mr = MeshResource(dp_resource="dp", ep_resource="ep") + # Drop any communicator a prior test class bootstrapped, then re-init. + ep_finalize() + with cls.mesh, global_shard_guard(cls.mr): + ep_bootstrap( + world_size=cls.num_procs, + rank=cls.rank, + num_experts=cls.num_experts, + max_tokens_per_rank=cls.max_tokens_per_rank, + recv_capacity_per_rank=cls.recv_capacity_per_rank, + hidden_dim=HIDDEN_DIM, + drop_on_overflow=True, + ) + cls.hk = EpLayerConfig(top_k=TOP_K, dispatch_output_per_expert_alignment=cls.ALIGN) + + @classmethod + def tearDownClass(cls): + # Leave a clean slate so another class can bootstrap after us. + ep_finalize() + + def _make_concentrated_inputs(self): + """All top-1 routes to expert 0; top-2 spread over the rest, so the rank + owning expert 0 overloads beyond recv_capacity_per_rank.""" + T_global = self.OVF_TOKENS_PER_DP_SHARD * self.dp + E = self.num_experts + topk_idx = np.empty((T_global, TOP_K), dtype=np.int32) + for t in range(T_global): + topk_idx[t, 0] = 0 + for k in range(1, TOP_K): + topk_idx[t, k] = 1 + (t % (E - 1)) + topk_idx = jnp.asarray(topk_idx) + topk_w = jnp.full((T_global, TOP_K), 1.0 / TOP_K, dtype=jnp.float32) + tokens = jnp.asarray( + np.linspace(0.1, 0.9, T_global * HIDDEN_DIM, dtype=np.float32).reshape( + T_global, HIDDEN_DIM + ), + dtype=jnp.bfloat16, + ) + return T_global, topk_idx, tokens, topk_w + + def test_overflow_drops_without_trap(self): + """Dispatch drops the overflow instead of trapping; total_recv_tokens + reports the pre-drop demand, which exceeds recv_capacity_per_rank.""" + _T, topk_idx, tokens, topk_w = self._make_concentrated_inputs() + dp_spec = PartitionSpec(("dp", "ep"), None) + ep_spec_3d = PartitionSpec(("dp", "ep"), None, None) + with self.mesh, global_shard_guard(self.mr): + + @jax.jit + def run(idx, toks, w): + idx = jax.lax.with_sharding_constraint(idx, NamedSharding(self.mesh, dp_spec)) + toks = jax.lax.with_sharding_constraint(toks, NamedSharding(self.mesh, dp_spec)) + w = jax.lax.with_sharding_constraint(w, NamedSharding(self.mesh, dp_spec)) + _tc, trt_prep, _hm = ep_prepare(self.hk, idx) + recv_t, _rw, _hm2, _tc2, trt_disp = ep_dispatch( + self.hk, idx, toks, w, self.recv_capacity_per_rank + ) + recv_t = jax.lax.with_sharding_constraint( + recv_t, NamedSharding(self.mesh, ep_spec_3d) + ) + return trt_prep, trt_disp, recv_t + + trt_prep, trt_disp, recv_t = run(topk_idx, tokens, topk_w) + recv_t.block_until_ready() + + # Reaching here means the dispatch dropped the overflow rather than + # trapping (an overflow trap is a device-side abort). + self.assertEqual(recv_t.shape[1], self.recv_capacity_per_rank) + # total_recv_tokens is sharded across processes; gather before host reads. + trt_prep = np.asarray(jmu.process_allgather(trt_prep, tiled=True)).reshape(-1) + trt_disp = np.asarray(jmu.process_allgather(trt_disp, tiled=True)).reshape(-1) + # prepare and dispatch see the same routing -> identical pre-drop totals. + np.testing.assert_array_equal(trt_prep, trt_disp) + # The rank owning expert 0 demands more than it can receive. + self.assertGreater(int(trt_prep.max()), self.recv_capacity_per_rank) + + +# ── EP domain grouping (single-process; runs under plain pytest) ───────────── + + +class TestEpDomainGrouping(unittest.TestCase): + """EP domains group ranks sharing all non-ep coords, so an orthogonal tp + axis splits the world into one EP domain per tp coordinate.""" + + def test_ep_tp_splits_domains(self): + # Gate on device count inside the test: calling jax.devices() at + # class-definition time would initialize the XLA backend before + # jax.distributed.initialize(). + if not is_devices_enough(8): + self.skipTest("requires 8 devices") + # ep=4, tp=2: tp must yield 2 EP domains, each a fixed tp coordinate. + mesh = Mesh(np.asarray(jax.devices()[:8]).reshape(4, 2), ("expert", "tensor")) + # Single host shares one process_index; inject row-major ranks to mimic + # one device per process. + order = {int(d.id): i for i, d in enumerate(mesh.devices.reshape(-1))} + d2r = lambda d: order[int(d.id)] + + domains = {} + for rank in range(8): + root, col, ndom = _ep_domain_for_rank(mesh, "expert", rank, device_to_rank=d2r) + self.assertEqual(ndom, 2) # every rank must agree on the domain count + domains.setdefault(root, {})[col] = rank + domains = {root: [m[c] for c in sorted(m)] for root, m in domains.items()} + + self.assertEqual(domains, {0: [0, 2, 4, 6], 1: [1, 3, 5, 7]}) + + # ── Entry point ────────────────────────────────────────────────────────────── @@ -732,12 +981,14 @@ def bwd_only(eo, toks, idx, w, g): ) loader = unittest.TestLoader() + test_cases = (TestEP, TestEPOverflowDrop, TestEpDomainGrouping) target = os.environ.get("TARGET_TEST") if target: name = target.split(".")[-1] - suite = loader.loadTestsFromName(name, TestEP) + cls = next((c for c in test_cases if hasattr(c, name)), TestEP) + suite = loader.loadTestsFromName(name, cls) else: - suite = loader.loadTestsFromTestCase(TestEP) + suite = unittest.TestSuite(loader.loadTestsFromTestCase(c) for c in test_cases) runner = unittest.TextTestRunner(verbosity=2) result = runner.run(suite) sys.exit(0 if result.wasSuccessful() else 1) diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py index d08765e184..015c73343a 100644 --- a/tests/jax/test_te_ep_moe.py +++ b/tests/jax/test_te_ep_moe.py @@ -31,11 +31,12 @@ on the block are pytest parametrize values rather than separate test classes: -* ``test_forward`` covers the forward across a curated set of - configurations (softmax/sigmoid scoring, optional non-zero - expert_bias). Each config asserts shape, dtype, finiteness and - numerical parity vs the reference in one run. -* ``test_backward`` mirrors that for gradients. +* ``test_forward`` covers BF16 and MXFP8 forward execution across a + curated set of configurations (softmax/sigmoid scoring, optional + non-zero expert_bias). Each config asserts shape, dtype, finiteness + and numerical parity vs the same BF16 reference in one run. +* ``test_backward`` mirrors that for gradients. BF16 and MXFP8 share + the full test body and differ only in the grouped-GEMM quantizer sets. * ``TestTeEpMoeAuxLoss`` covers the second return value end-to-end (returned + parity + aux-only grad propagates to gate + combined main+aux grads stay finite) in two consolidated tests. @@ -118,8 +119,14 @@ def _read_mp_options(): ) from transformer_engine.jax.flax import _MoEBlock as MoEBlock -from transformer_engine.jax.moe import _ALIGN_SIZE, moe, record_ep_bootstrap_signature_for_moe +from transformer_engine.jax.moe import ( + _ALIGN_SIZE, + get_moe_recv_capacity_per_rank, + moe, + record_ep_bootstrap_signature_for_moe, +) from transformer_engine.jax.ep import ep_bootstrap +from transformer_engine.common.recipe import MXFP8BlockScaling from transformer_engine.jax.sharding import MeshResource, global_shard_guard @@ -140,7 +147,7 @@ def _read_mp_options(): ("exp", EP_AXIS), ("embed", FSDP_AXIS), ("mlp", None), - ("batch", (EP_AXIS, FSDP_AXIS)), + ("batch", (FSDP_AXIS, EP_AXIS)), ) # Small shapes so the parity tests stay tight on bf16. The block still @@ -148,31 +155,31 @@ def _read_mp_options(): DTYPE = jnp.bfloat16 BATCH = EP_SIZE * FSDP_SIZE * 2 # 8 on 4-GPU, 16 on 8-GPU SEQ = 32 -HIDDEN = 64 +HIDDEN = 128 INTER = 128 NUM_EXPERTS = 8 TOPK = 2 -# bf16 grouped_gemm + softmax-topk + ep all-to-all stack drifts ~1e-1 vs a -# fp32 numpy reference. Keep these tight enough to catch real bugs but -# loose enough to absorb expected bf16 rounding. -FWD_ATOL = 5e-2 -FWD_RTOL = 5e-2 -GRAD_FFN_ATOL = 1e-1 -GRAD_FFN_RTOL = 1e-1 -GRAD_GATE_ATOL = 5e-1 -GRAD_GATE_RTOL = 5e-1 - -# Two TE EP runs that should be bitwise-equal modulo XLA fusion order -# (slot alignment rounding, etc.). -TE_TO_TE_ATOL = 5e-3 -TE_TO_TE_RTOL = 5e-3 +# MXFP8 grouped GEMMs have measurably more quantization drift than BF16. +# These bounds are rounded slightly above the worst error observed across +# the forward and backward configuration matrix. +FWD_TOLERANCE = { + "bf16": {"atol": 5e-4, "rtol": 5e-4}, + "mxfp8": {"atol": 7e-3, "rtol": 7e-3}, +} +GRAD_FFN_TOLERANCE = { + "bf16": {"atol": 1e-7, "rtol": 1e-7}, + "mxfp8": {"atol": 1.3e-6, "rtol": 1.3e-6}, +} +GRAD_GATE_TOLERANCE = { + "bf16": {"atol": 7e-8, "rtol": 7e-8}, + "mxfp8": {"atol": 8e-7, "rtol": 8e-7}, +} # Aux loss is computed in float32 from the SAME logits as the routing # path. Numerical drift between TE-EP and the reference is dominated by # the bf16-rounded softmax inside the topk kernel. -AUX_ATOL = 1e-3 -AUX_RTOL = 1e-3 +AUX_TOLERANCE = {"atol": 1e-6, "rtol": 1e-6} # ----------------------------------------------------------------------------- @@ -180,28 +187,6 @@ def _read_mp_options(): # ----------------------------------------------------------------------------- -def _compute_worst_case_recv_pr(): - """Per-rank recv buffer the bootstrap must reserve. - - NCCL EP HT expert-major uses one flat recv buffer with variable - per-expert zones. Each non-empty expert zone is padded to - ``_ALIGN_SIZE`` slots, so the reserve must cover the worst-case - total assignments plus independent per-zone padding. - """ - num_procs = jax.device_count() - num_local_experts = NUM_EXPERTS // EP_SIZE - max_tokens_per_rank = (BATCH // num_procs) * SEQ - tokens_per_ep_group = EP_SIZE * max_tokens_per_rank - max_local_assignments = tokens_per_ep_group * min(TOPK, num_local_experts) - max_nonempty_experts = min(num_local_experts, max_local_assignments) - padded_total_bound = max_local_assignments + (_ALIGN_SIZE - 1) * max_nonempty_experts - aligned_total_bound = ((padded_total_bound + _ALIGN_SIZE - 1) // _ALIGN_SIZE) * _ALIGN_SIZE - per_expert_bound = ( - num_local_experts * ((tokens_per_ep_group + _ALIGN_SIZE - 1) // _ALIGN_SIZE) * _ALIGN_SIZE - ) - return min(per_expert_bound, aligned_total_bound) - - @pytest.fixture(scope="module") def mesh(): if jax.device_count() < NUM_DEVICES_REQUIRED: @@ -217,7 +202,15 @@ def mesh(): num_procs = jax.process_count() max_tokens_per_rank = (BATCH // num_procs) * SEQ - recv_capacity_per_rank = _compute_worst_case_recv_pr() + # Worst-case recv capacity per rank + # TODO(jberchtold) support configurations other than worst-case by refactoring tests + # but if possible avoid bootstrap/teardown for each test + recv_capacity_per_rank = get_moe_recv_capacity_per_rank( + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + max_tokens_per_rank=max_tokens_per_rank, + ep_size=EP_SIZE, + ) # Eager bootstrap: ep_bootstrap does a host-side NCCL UID allgather # and cannot run from inside jax.jit. Sized to the worst-case recv_pr @@ -275,8 +268,7 @@ def mesh(): def _pure_jax_moe_reference( x, gate_kernel, - wi_0, - wi_1, + wi, wo, expert_bias=None, *, @@ -317,6 +309,7 @@ def _pure_jax_moe_reference( # FFN. ``apply_topk_weights_early`` is a fusion knob that doesn't # change the math (wo is linear), so the reference is identical for # both placements. + wi_0, wi_1 = jnp.split(wi, 2, axis=-1) layer_w0 = jnp.einsum("th,ehm->tem", x_2d, wi_0) layer_w1 = jnp.einsum("th,ehm->tem", x_2d, wi_1) # Activation runs in x.dtype (typically bf16) to mirror the impl -- @@ -364,6 +357,8 @@ def _make_block( use_expert_routing_bias=False, score_function="softmax", expert_bias_init=None, + input_axes=("batch", None, None), + quantization_recipe=None, ): kwargs = dict( num_experts=NUM_EXPERTS, @@ -375,6 +370,8 @@ def _make_block( use_expert_routing_bias=use_expert_routing_bias, score_function=score_function, dtype=DTYPE, + input_axes=input_axes, + quantization_recipe=quantization_recipe, ) # Custom expert_bias_init lets tests inject a non-zero expert_bias without # poking variables['params'] post-init. @@ -429,12 +426,19 @@ def _init_apply(block, mesh, x, key): x_sh = _shard_inputs(x, mesh) variables = jax.jit(block.init)(key, x_sh) jax.block_until_ready(jax.tree_util.tree_leaves(variables)[0]) - output, aux = jax.jit(block.apply)(variables, x_sh) + output, aux, _trt = jax.jit(block.apply)(variables, x_sh) jax.block_until_ready(output) return variables, output, aux -def _grad_step(block, variables, mesh, x, *, include_aux=False): +def _grad_step( + block, + variables, + mesh, + x, + *, + include_aux=False, +): """Run jax.grad of mean(out^2) [+ aux if include_aux] vs (params, x). Returns ``(grads_variables, grad_x)`` so callers can check both the @@ -445,7 +449,7 @@ def _grad_step(block, variables, mesh, x, *, include_aux=False): x_sh = _shard_inputs(x, mesh) def loss_fn(variables, x): - output, aux = block.apply(variables, x) + output, aux, _trt = block.apply(variables, x) loss = jnp.mean(output.astype(jnp.float32) ** 2) if include_aux and aux is not None: loss = loss + aux.astype(jnp.float32) @@ -464,7 +468,7 @@ def _grad_aux_only(block, variables, mesh, x): x_sh = _shard_inputs(x, mesh) def aux_only(variables, x): - _, aux = block.apply(variables, x) + _, aux, _trt = block.apply(variables, x) return aux.astype(jnp.float32) grads = jax.jit(jax.grad(aux_only))(variables, x_sh) @@ -501,6 +505,13 @@ def _make_inputs(key): return jax.random.normal(key, (BATCH, SEQ, HIDDEN), dtype=DTYPE) +def _quantization_recipe(quantization): + if quantization == "bf16": + return None + assert quantization == "mxfp8" + return MXFP8BlockScaling() + + # ----------------------------------------------------------------------------- # Tests # ----------------------------------------------------------------------------- @@ -544,6 +555,13 @@ def _make_inputs(key): ), ] +_QUANTIZATION_CASES = [ + pytest.param("bf16", id="bf16"), +] + +if get_device_compute_capability(0) >= 100: + _QUANTIZATION_CASES.append(pytest.param("mxfp8", id="mxfp8")) + def _reference_kwargs_from_config(config, params_np): """Pick out the reference-relevant pieces of a parametrize config.""" @@ -562,8 +580,9 @@ class TestTeEpMoeForward: finiteness AND numerical parity vs the pure-JAX reference.""" @pytest.mark.parametrize("config", _CONFIGS) - def test_forward(self, mesh, config): - block = _make_block(**config) + @pytest.mark.parametrize("quantization", _QUANTIZATION_CASES) + def test_forward(self, mesh, config, quantization): + block = _make_block(**config, quantization_recipe=_quantization_recipe(quantization)) x = _make_inputs(jax.random.PRNGKey(0)) variables, output, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(1)) @@ -582,8 +601,7 @@ def test_forward(self, mesh, config): out_ref, _ = _pure_jax_moe_reference( jnp.asarray(x_np), jnp.asarray(params_np["gate_kernel"]), - jnp.asarray(params_np["wi_0"]), - jnp.asarray(params_np["wi_1"]), + jnp.asarray(params_np["wi"]), jnp.asarray(params_np["wo"]), num_experts=NUM_EXPERTS, num_experts_per_tok=TOPK, @@ -592,9 +610,8 @@ def test_forward(self, mesh, config): np.testing.assert_allclose( out_te_np.astype(np.float32), np.asarray(jax.device_get(out_ref)).astype(np.float32), - atol=FWD_ATOL, - rtol=FWD_RTOL, - err_msg=f"forward parity breach for config={config}", + **FWD_TOLERANCE[quantization], + err_msg=f"forward parity breach for config={config}, quantization={quantization}", ) @@ -603,8 +620,9 @@ class TestTeEpMoeBackward: grads finite, non-zero AND parity vs the pure-JAX reference.""" @pytest.mark.parametrize("config", _CONFIGS) - def test_backward(self, mesh, config): - block = _make_block(**config) + @pytest.mark.parametrize("quantization", _QUANTIZATION_CASES) + def test_backward(self, mesh, config, quantization): + block = _make_block(**config, quantization_recipe=_quantization_recipe(quantization)) x = _make_inputs(jax.random.PRNGKey(2)) variables, _, _ = _init_apply(block, mesh, x, jax.random.PRNGKey(3)) grads_te, grad_x_te = _grad_step(block, variables, mesh, x) @@ -621,8 +639,7 @@ def loss_fn(params, x): out, _ = _pure_jax_moe_reference( x, params["gate_kernel"], - params["wi_0"], - params["wi_1"], + params["wi"], params["wo"], ref_expert_bias, num_experts=NUM_EXPERTS, @@ -638,22 +655,19 @@ def loss_fn(params, x): grads_ref_np = {k: np.asarray(jax.device_get(v)) for k, v in grads_ref.items()} grad_x_ref_np = np.asarray(jax.device_get(grad_x_ref)) - for name in ("gate_kernel", "wi_0", "wi_1", "wo"): + for name in ("gate_kernel", "wi", "wo"): # Per-tensor: finite + non-zero + parity in one pass. g_te = _to_global_numpy(_unwrap(grads_te["params"][name]), mesh) assert np.all(np.isfinite(g_te)), f"{name} grad has NaN/Inf [config={config}]" assert np.any(g_te != 0.0), f"{name} grad identically zero [config={config}]" - atol, rtol = ( - (GRAD_GATE_ATOL, GRAD_GATE_RTOL) - if name == "gate_kernel" - else (GRAD_FFN_ATOL, GRAD_FFN_RTOL) - ) + tolerances = GRAD_GATE_TOLERANCE if name == "gate_kernel" else GRAD_FFN_TOLERANCE np.testing.assert_allclose( g_te.astype(np.float32), grads_ref_np[name].astype(np.float32), - atol=atol, - rtol=rtol, - err_msg=f"grad parity breach on {name} [config={config}]", + **tolerances[quantization], + err_msg=( + f"grad parity breach on {name} [config={config}, quantization={quantization}]" + ), ) # d_x: the gradient propagated back to the previous layer. Checks @@ -673,9 +687,8 @@ def loss_fn(params, x): np.testing.assert_allclose( grad_x_te_np.astype(np.float32), grad_x_ref_np.astype(np.float32), - atol=GRAD_FFN_ATOL, - rtol=GRAD_FFN_RTOL, - err_msg=f"d_x parity breach [config={config}]", + **GRAD_FFN_TOLERANCE[quantization], + err_msg=f"d_x parity breach [config={config}, quantization={quantization}]", ) @@ -708,8 +721,7 @@ def test_aux_loss(self, mesh): _, aux_ref = _pure_jax_moe_reference( jnp.asarray(x_np), jnp.asarray(params_np["gate_kernel"]), - jnp.asarray(params_np["wi_0"]), - jnp.asarray(params_np["wi_1"]), + jnp.asarray(params_np["wi"]), jnp.asarray(params_np["wo"]), num_experts=NUM_EXPERTS, num_experts_per_tok=TOPK, @@ -718,8 +730,7 @@ def test_aux_loss(self, mesh): np.testing.assert_allclose( float(aux_np), float(jax.device_get(aux_ref)), - atol=AUX_ATOL, - rtol=AUX_RTOL, + **AUX_TOLERANCE, ) # Aux-only bwd must propagate to gate_kernel — proves the @@ -739,7 +750,7 @@ def test_combined_loss_grads(self, mesh): x = _make_inputs(jax.random.PRNGKey(22)) variables, _, _ = _init_apply(block, mesh, x, jax.random.PRNGKey(23)) grads, _ = _grad_step(block, variables, mesh, x, include_aux=True) - for name in ("gate_kernel", "wi_0", "wi_1", "wo"): + for name in ("gate_kernel", "wi", "wo"): g_local = np.asarray(jax.device_get(_unwrap(grads["params"][name]).addressable_data(0))) assert np.all(np.isfinite(g_local)), f"{name} grad NaN/Inf under main+aux" assert np.any(g_local != 0.0), f"{name} grad zero under main+aux" diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py index 8e7ffb8e14..ef2363cf81 100644 --- a/tests/pytorch/attention/run_attention_with_cp.py +++ b/tests/pytorch/attention/run_attention_with_cp.py @@ -25,7 +25,10 @@ # Executed as a script, so sibling imports rely on the interpreter putting this file's # directory on sys.path -- which safe-path mode (PYTHONSAFEPATH, python -P) disables. sys.path.append(os.path.dirname(os.path.realpath(__file__))) -from test_attention_with_cp import model_configs_flash_attn, model_configs_fused_attn +from test_attention_with_cp import ( + model_configs_flash_attn, + model_configs_fused_attn, +) from transformer_engine.pytorch import ( autocast, DotProductAttention, @@ -226,6 +229,11 @@ def run_dpa_with_cp( logging.root.setLevel(log_level) # When is_training is False, gradient outputs are None. is_training = is_training == "True" + pad_between_seqs = None + if qkv_format == "thd": + # Keep this in sync with generate_input_shapes so DPA gets the explicit + # padding state without a GPU-to-CPU sync. + pad_between_seqs = kernel_backend == "FusedAttention" or fa_pad_between_seqs == "True" # set up environment variables and config if deterministic == "True": @@ -248,7 +256,10 @@ def run_dpa_with_cp( config = copy.deepcopy(model_configs_flash_attn[model]) if kernel_backend == "FusedAttention": os.environ["NVTE_FUSED_ATTN"] = "1" - config = copy.deepcopy(model_configs_fused_attn[model]) + if model in model_configs_fused_attn: + config = copy.deepcopy(model_configs_fused_attn[model]) + else: + assert False, f"{model=} is not a known FusedAttention CP config!" assert config.attn_mask_type in [ "causal", "no_mask", @@ -428,13 +439,7 @@ def run_dpa_with_cp( cu_seqlens_kv=cu_seqlens_kv, cu_seqlens_q_padded=cu_seqlens_q_padded, cu_seqlens_kv_padded=cu_seqlens_kv_padded, - # Test runner sets cu_seqlens_q == cu_seqlens_q_padded for the - # FlashAttention path, i.e. no inter-sequence padding. Declare this - # explicitly so the sync-free auto-detect (which conservatively - # picks True when padded cu_seqlens are present) does not disable FA. - pad_between_seqs=( - (kernel_backend != "FlashAttention") if qkv_format == "thd" else None - ), + pad_between_seqs=pad_between_seqs, fp8_output=fp8_mha, ) if config.return_max_logit: @@ -552,12 +557,7 @@ def run_dpa_with_cp( cu_seqlens_kv=cu_seqlens_kv, cu_seqlens_q_padded=cu_seqlens_q_padded, cu_seqlens_kv_padded=cu_seqlens_kv_padded, - # See note above (non-CP branch): same explicit declaration so - # FlashAttention isn't disabled by the conservative sync-free - # auto-detect when this test path constructs no inter-seq padding. - pad_between_seqs=( - (kernel_backend != "FlashAttention") if qkv_format == "thd" else None - ), + pad_between_seqs=pad_between_seqs, fp8_output=fp8_mha, ) if config.return_max_logit: diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 1726f9a4c3..f973faffd9 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -7,8 +7,10 @@ import os import sys import pathlib +import copy from typing import Any, Dict, Tuple, Union +from packaging.version import Version as PkgVersion import pytest import torch @@ -32,6 +34,7 @@ ) from transformer_engine.pytorch.attention.dot_product_attention.utils import ( FlashAttentionUtils, + _get_supported_versions, check_set_window_size, ) from transformer_engine.pytorch.attention import RotaryPositionEmbedding @@ -114,6 +117,40 @@ def reset_attn_backend(): "NVTE_CK_USES_FWD_V3", "NVTE_CK_USES_BWD_V3", "NVTE_FP8_DPA_BWD"]) yield + +@pytest.mark.parametrize( + "version,expected", + ( + ("2.1.0", False), + ("2.1.1", True), + ("2.8.3", True), + ("2.8.3+local_version", True), + ("2.8.3.post1", True), + ("2.8.4", False), + ("2.8.4+local_version", False), + ("2.9.0", False), + ), +) +def test_flash_attention_version_support(version, expected): + """Test the supported Flash Attention v2 version range.""" + assert ( + FlashAttentionUtils.is_version_supported( + PkgVersion(version), FlashAttentionUtils.version_required + ) + is expected + ) + + +def test_flash_attention_supported_version_message(): + """Test that the supported version range describes an exclusive upper bound.""" + assert ( + _get_supported_versions( + FlashAttentionUtils.version_required, FlashAttentionUtils.max_version + ) + == ">= 2.1.1, < 2.8.4" + ) + + # Define F16 data types to test param_types = [torch.float16] if is_bf16_available(): @@ -224,6 +261,8 @@ def test_dot_product_attention( qkv_layout, swa, pad_between_seqs, + declarative_packed=False, + is_training=True, ): """Test DotProductAttention module""" @@ -231,7 +270,7 @@ def test_dot_product_attention( tols = dict(atol=1e-3, rtol=1e-3) if dtype == torch.bfloat16: tols = dict(atol=1.5e-2, rtol=1.5e-2) - config = model_configs[model] + config = copy.deepcopy(model_configs[model]) is_mla = config.head_dim_qk != config.head_dim_v is_mqa_gqa = config.num_heads != config.num_gqa_groups if qkv_layout is None: @@ -241,6 +280,8 @@ def test_dot_product_attention( qkv_layout = "bshd_bs2hd" if not is_mla and not is_mqa_gqa else "bshd_bshd_bshd" if "3" in qkv_layout and config.attn_type == "cross": pytest.skip("No need to test this layout for cross attention") + if declarative_packed and not any(c.isdigit() for c in qkv_layout): + pytest.skip("Declarative packed inputs only apply to packed qkv layouts.") if config.window_size == (-1, -1) and swa: config.window_size = [2, 2] @@ -256,10 +297,9 @@ def test_dot_product_attention( # Get backends # For 111s, dbias calculation is not supported as of cuDNN 9.18, hence, test fwd only for 111s. - # For all other shapes test fwd+bwd - is_training = True + # For all other shapes test fwd+bwd unless the caller requests fwd-only coverage. # TODO(KshitijLakhani): Set is_training to True for all cases once cuDNN supports dbias for 111s. - if config.bias_shape == "111s": + if is_training and config.bias_shape == "111s": is_training = False logging.info( "Setting is_training to False as cuDNN does not support dbias for" @@ -327,7 +367,61 @@ def test_dot_product_attention( # FusedAttention backend if fused_attn_supported: - if len(fused_attn_backends) == 1: + if IS_HIP_EXTENSION: + # ROCm exercises multiple fused-attn backends (CK V2/V3, AOTriton) when + # more than one is available, comparing them against each other below. + if len(fused_attn_backends) == 1: + fused_attn_fwd, fused_max_logit, fused_attn_bwd = _run_dot_product_attention( + dtype, + config, + "FusedAttention", + ckpt_attn, + qkv_layout, + pad_between_seqs, + is_training, + ) + if len(fused_attn_backends) == 2: + os.environ["NVTE_FUSED_ATTN_BACKEND"] = "0" + os.environ["NVTE_FUSED_ATTN_CK"] = "0" + os.environ["NVTE_FUSED_ATTN_AOTRITON"] = "1" + fused_attn_fwd, _, fused_attn_bwd = _run_dot_product_attention( + dtype, + config, + "FusedAttention", + ckpt_attn, + qkv_layout, + pad_between_seqs, + is_training, + ) + os.environ["NVTE_FUSED_ATTN_BACKEND"] = "1" + os.environ["NVTE_FUSED_ATTN_CK"] = "1" + os.environ["NVTE_FUSED_ATTN_AOTRITON"] = "0" + os.environ["NVTE_CK_USES_FWD_V3"] = "1" + os.environ["NVTE_CK_USES_BWD_V3"] = "1" + fused_attn_fwd_1, _, fused_attn_bwd_1 = _run_dot_product_attention( + dtype, + config, + "FusedAttention", + ckpt_attn, + qkv_layout, + pad_between_seqs, + is_training, + ) + if has_ck_backend: + os.environ["NVTE_FUSED_ATTN_CK"] = "1" + os.environ["NVTE_FUSED_ATTN_AOTRITON"] = "0" + os.environ["NVTE_CK_USES_FWD_V3"] = "0" + os.environ["NVTE_CK_USES_BWD_V3"] = "0" + fused_attn_fwd_2, _, fused_attn_bwd_2 = _run_dot_product_attention( + dtype, + config, + "FusedAttention", + ckpt_attn, + qkv_layout, + pad_between_seqs, + is_training, + ) + else: fused_attn_fwd, fused_max_logit, fused_attn_bwd = _run_dot_product_attention( dtype, config, @@ -336,47 +430,7 @@ def test_dot_product_attention( qkv_layout, pad_between_seqs, is_training, - ) - if len(fused_attn_backends) == 2: - os.environ["NVTE_FUSED_ATTN_BACKEND"] = "0" - os.environ["NVTE_FUSED_ATTN_CK"] = "0" - os.environ["NVTE_FUSED_ATTN_AOTRITON"] = "1" - fused_attn_fwd, _, fused_attn_bwd = _run_dot_product_attention( - dtype, - config, - "FusedAttention", - ckpt_attn, - qkv_layout, - pad_between_seqs, - is_training, - ) - os.environ["NVTE_FUSED_ATTN_BACKEND"] = "1" - os.environ["NVTE_FUSED_ATTN_CK"] = "1" - os.environ["NVTE_FUSED_ATTN_AOTRITON"] = "0" - os.environ["NVTE_CK_USES_FWD_V3"] = "1" - os.environ["NVTE_CK_USES_BWD_V3"] = "1" - fused_attn_fwd_1, _, fused_attn_bwd_1 = _run_dot_product_attention( - dtype, - config, - "FusedAttention", - ckpt_attn, - qkv_layout, - pad_between_seqs, - is_training, - ) - if has_ck_backend: - os.environ["NVTE_FUSED_ATTN_CK"] = "1" - os.environ["NVTE_FUSED_ATTN_AOTRITON"] = "0" - os.environ["NVTE_CK_USES_FWD_V3"] = "0" - os.environ["NVTE_CK_USES_BWD_V3"] = "0" - fused_attn_fwd_2, _, fused_attn_bwd_2 = _run_dot_product_attention( - dtype, - config, - "FusedAttention", - ckpt_attn, - qkv_layout, - pad_between_seqs, - is_training, + declarative_packed=declarative_packed, ) # FlashAttention backend @@ -389,6 +443,7 @@ def test_dot_product_attention( qkv_layout, pad_between_seqs, is_training, + declarative_packed=declarative_packed, ) # Compare results @@ -501,9 +556,18 @@ def test_dpa_num_splits(dtype, model_configs, model): } -@pytest.mark.skipif( - not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." +fa4_enabled = bool(int(os.getenv("NVTE_FLASH_ATTN", "1"))) and bool( + int(os.getenv("NVTE_FLASH_ATTN_V4", "1")) ) +requires_fa4 = pytest.mark.skipif( + not fa4_enabled + or not FlashAttentionUtils.v4_is_installed + or device_compute_capability < (9, 0), + reason="Enabled Flash-attn v4 and compute capability >= SM90 are required.", +) + + +@requires_fa4 @pytest.mark.parametrize("dtype", param_types_lean) @pytest.mark.parametrize("model_configs", [model_configs_fa4_base]) @pytest.mark.parametrize("model", model_configs_fa4_base.keys()) @@ -522,9 +586,7 @@ def test_dpa_fa4_base(dtype, model_configs, model): } -@pytest.mark.skipif( - not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." -) +@requires_fa4 @pytest.mark.skipif( device_compute_capability not in ((10, 0), (10, 3)), reason="FA4 head_dim=256 dedicated kernel is SM100/103-only.", @@ -534,7 +596,59 @@ def test_dpa_fa4_base(dtype, model_configs, model): @pytest.mark.parametrize("model", model_configs_fa4_hdim256.keys()) def test_dpa_fa4_hdim256(dtype, model_configs, model): """Test DotProductAttention with FA4: head_dim=256 dedicated kernel on SM100""" - test_dot_product_attention(dtype, model_configs, model, False, None, False, False) + # Keep this FA4 D=256 test forward-only. Before cuDNN D=256 backward support, + # the generic helper took this path implicitly because fused-attn training was unavailable. + test_dot_product_attention( + dtype, model_configs, model, False, None, False, False, is_training=False + ) + + +# cuDNN FusedAttention D=256 bprop is supported on sm10x by the dedicated deterministic +# SDPA bprop kernel. BSHD support starts with cuDNN FE 1.24 / BE 9.23; THD support starts +# with cuDNN FE 1.26 / BE 9.25. The kernel supports d_qk == d_v == 256 only, vanilla softmax only, +# no dropout, no ALiBi, and (for non-causal masks) full-window attention only. +model_configs_d256 = { + # test: ModelConfig(b, sq, hq, dqk) -> head_dim_v defaults to head_dim_qk (256) + "d256_no_mask": ModelConfig(2, 512, 16, 256), + "d256_padding": ModelConfig(2, 512, 16, 256, attn_mask_type="padding"), + # SWA is allowed only together with a causal mask on the D=256 bprop kernel. + "d256_causal_swa": ModelConfig(2, 1024, 16, 256, attn_mask_type="causal", window_size=(128, 0)), + # GQA variant (num_gqa_groups < num_heads). + "d256_padding_causal_gqa": ModelConfig( + 2, 1024, 16, 256, num_gqa_groups=4, attn_mask_type="padding_causal" + ), +} + + +@pytest.mark.skipif( + device_compute_capability not in ((10, 0), (10, 3)), + reason="cuDNN FusedAttention head_dim=256 backward is Blackwell server (SM100/SM103) only.", +) +@pytest.mark.parametrize("dtype", param_types) +@pytest.mark.parametrize("model_configs", [model_configs_d256]) +@pytest.mark.parametrize("model", model_configs_d256.keys()) +@pytest.mark.parametrize( + "qkv_layout", + [ + pytest.param( + "bshd_bs2hd", + marks=pytest.mark.skipif( + get_cudnn_version() < (9, 23, 0), + reason="cuDNN 9.23+ is required for BSHD D=256 fused-attn backward.", + ), + ), + pytest.param( + "thd_t2hd", + marks=pytest.mark.skipif( + get_cudnn_version() < (9, 25, 0), + reason="cuDNN 9.25+ is required for THD D=256 fused-attn backward.", + ), + ), + ], +) +def test_dpa_d256(dtype, model_configs, model, qkv_layout): + """Test DotProductAttention with head_dim=256 backward on Blackwell""" + test_dot_product_attention(dtype, model_configs, model, False, qkv_layout, False, False) model_configs_fa4_mla = { @@ -550,9 +664,7 @@ def test_dpa_fa4_hdim256(dtype, model_configs, model): } -@pytest.mark.skipif( - not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." -) +@requires_fa4 @pytest.mark.parametrize("dtype", param_types_lean) @pytest.mark.parametrize("model_configs", [model_configs_fa4_mla]) @pytest.mark.parametrize("model", model_configs_fa4_mla.keys()) @@ -574,9 +686,7 @@ def test_dpa_fa4_mla(dtype, model_configs, model): } -@pytest.mark.skipif( - not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." -) +@requires_fa4 @pytest.mark.parametrize("dtype", param_types_lean) @pytest.mark.parametrize("model_configs", [model_configs_fa4_swa]) @pytest.mark.parametrize("model", model_configs_fa4_swa.keys()) @@ -597,9 +707,7 @@ def test_dpa_fa4_sliding_window(dtype, model_configs, model, qkv_layout): } -@pytest.mark.skipif( - not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." -) +@requires_fa4 @pytest.mark.parametrize("dtype", param_types_lean) @pytest.mark.parametrize("model_configs", [model_configs_fa4_varlen]) @pytest.mark.parametrize("model", model_configs_fa4_varlen.keys()) @@ -622,9 +730,7 @@ def test_dpa_fa4_varlen(dtype, model_configs, model, qkv_layout): } -@pytest.mark.skipif( - not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." -) +@requires_fa4 @pytest.mark.parametrize("dtype", param_types_lean) @pytest.mark.parametrize("model_configs", [model_configs_fa4_mask]) @pytest.mark.parametrize("model", model_configs_fa4_mask.keys()) @@ -1093,6 +1199,26 @@ def test_dpa_qkv_layout(dtype, model_configs, model, qkv_layout): test_dot_product_attention(dtype, model_configs, model, False, qkv_layout, False, False) +qkv_layouts_packed = [l for l in qkv_layouts if any(c.isdigit() for c in l)] + + +@pytest.mark.skipif(get_cudnn_version() < (8, 9, 5), reason="cuDNN 8.9.5+ is required.") +@pytest.mark.parametrize("dtype", param_types_lean) +@pytest.mark.parametrize("model_configs", [model_configs_layout]) +@pytest.mark.parametrize("model", ["layout_1_1", "layout_1_2"]) +@pytest.mark.parametrize("qkv_layout", qkv_layouts_packed) +def test_dpa_qkv_layout_declarative(dtype, model_configs, model, qkv_layout): + """Declarative packed inputs: the packed buffer is passed to + DotProductAttention via qkv_layer/kv_layer (declared layout, gradients read + off the packed buffer) instead of q/k/v views + pointer-based detection. + Layout coverage is complete; the model-config dimension is trimmed to one + self-attention and one cross-attention config, since past the input + handling the backend code is identical to test_dpa_qkv_layout.""" + test_dot_product_attention( + dtype, model_configs, model, False, qkv_layout, False, False, declarative_packed=True + ) + + qkv_layouts_thd = ["t3hd", "th3d", "thd_t2hd", "thd_th2d", "thd_thd_thd"] model_configs_layout_thd = { # test: ModelConfig(b, sq, hq, dqk) @@ -1160,7 +1286,9 @@ def test_dpa_qkv_layout(dtype, model_configs, model, qkv_layout): @pytest.mark.parametrize("model", model_configs_layout_thd.keys()) @pytest.mark.parametrize("qkv_layout", qkv_layouts_thd) @pytest.mark.parametrize("pad_between_seqs", [True, False]) -def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout, pad_between_seqs): +def test_dpa_qkv_layout_thd( + dtype, model_configs, model, qkv_layout, pad_between_seqs=True, declarative_packed=False +): """Test DotProductAttention module with different QKV layouts""" config = model_configs[model] if config.num_heads != config.num_gqa_groups and "3" in qkv_layout: @@ -1168,7 +1296,14 @@ def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout, pad_between if (pad_between_seqs==False and get_cudnn_version() < (9, 3, 0)): pytest.skip("cuDNN 9.3.0+ is required to run pad_between_seqs = False"); test_dot_product_attention( - dtype, model_configs, model, False, qkv_layout, False, pad_between_seqs + dtype, + model_configs, + model, + False, + qkv_layout, + False, + pad_between_seqs, + declarative_packed=declarative_packed, ) @pytest.mark.skipif(not IS_HIP_EXTENSION, reason="ROCm TE specific pytests.") @@ -1194,10 +1329,32 @@ def find_factors(x): if config.num_heads != config.num_gqa_groups and "3" in qkv_layout: continue test_dot_product_attention( - dtype, model_configs, model, False, qkv_layout, False, pad_between_seqs + dtype, + model_configs, + model, + False, + qkv_layout, + False, + pad_between_seqs, ) +qkv_layouts_thd_packed = [l for l in qkv_layouts_thd if any(c.isdigit() for c in l)] + + +@pytest.mark.skipif(get_cudnn_version() < (9, 0, 0), reason="cuDNN 9.0.0+ is required.") +@pytest.mark.skipif( + get_device_compute_capability() < (9, 0), reason="THD is only supported on Hopper+." +) +@pytest.mark.parametrize("dtype", param_types_lean) +@pytest.mark.parametrize("model_configs", [model_configs_layout_thd]) +@pytest.mark.parametrize("model", ["layout_0_0"]) +@pytest.mark.parametrize("qkv_layout", qkv_layouts_thd_packed) +def test_dpa_qkv_layout_thd_declarative(dtype, model_configs, model, qkv_layout): + """Declarative packed thd inputs, see test_dpa_qkv_layout_declarative.""" + test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout, declarative_packed=True) + + def _run_dot_product_attention( dtype: torch.dtype, config: ModelConfig, @@ -1206,8 +1363,14 @@ def _run_dot_product_attention( qkv_layout: str, pad_between_seqs: bool, is_training: bool, + declarative_packed: bool = False, ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: - """Run DotProductAttention module with one forward pass and one backward pass""" + """Run DotProductAttention module with one forward pass and one backward pass. + + With declarative_packed=True (packed qkv_layout only), the packed buffer is + passed to DotProductAttention directly via qkv_layer/kv_layer instead of + slicing it into q/k/v views, and input gradients are read off the packed + buffer itself.""" # Set RNG and environment varables reset_rng_states() os.environ["NVTE_FLASH_ATTN"] = "0" @@ -1404,6 +1567,12 @@ def _run_dot_product_attention( tensor_count = int(l) split_dim = dim break + if declarative_packed and split_dim != 0: + # The packed buffer is the autograd leaf; q/k/v below are non-leaf + # views of it, and DPA receives the buffer via qkv_layer/kv_layer. + tensor.requires_grad_() + packed_tensor = tensor + packed_interleave_dim = split_dim - tensor.dim() tensors = torch.split(tensor, 1, dim=split_dim) if split_dim != 0 else [tensor] tensors_orig = ( torch.split(tensor_orig, 1, dim=split_dim) if split_dim != 0 else [tensor_orig] @@ -1416,8 +1585,10 @@ def _run_dot_product_attention( inp.append(tensors[j]) inp_orig.append(tensors_orig[j]) for i in range(3): - inp[i].requires_grad = True - inp_orig[i].requires_grad = True + if inp[i].is_leaf: + inp[i].requires_grad = True + if inp_orig[i].is_leaf: + inp_orig[i].requires_grad = True # Create output gradient qkv_format_kv = "_".join(qkv_format) @@ -1499,10 +1670,21 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: k = inp[1] v = inp[2] d_out = out_grad + packed_kwargs = {} + if declarative_packed: + assert backend in ["FusedAttention", "FlashAttention"] + packed_kwargs["qkv_interleave_dim"] = packed_interleave_dim + if len(qkv_layout.split("_")) == 1: + packed_kwargs["qkv_layer"] = packed_tensor + q, k, v = None, None, None + else: + packed_kwargs["kv_layer"] = packed_tensor + k, v = None, None out = block( q, k, v, + **packed_kwargs, window_size=config.window_size, attention_mask=attention_mask, qkv_format=qkv_format, @@ -1532,15 +1714,33 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: if is_training: out.backward(d_out) + q_grad, k_grad, v_grad = None, None, None + if is_training: + if declarative_packed: + # Input gradients live on the packed buffer; slice them back out so + # the cross-backend comparisons below stay uniform with the + # separate-q/k/v path. + assert ( + packed_tensor.grad is not None and packed_tensor.grad.shape == packed_tensor.shape + ) + packed_grads = [ + packed_tensor.grad.select(packed_interleave_dim, j) + for j in range(packed_tensor.shape[packed_interleave_dim]) + ] + if len(qkv_layout.split("_")) == 1: + q_grad, k_grad, v_grad = packed_grads + else: + q_grad = q.grad + k_grad, v_grad = packed_grads + else: + q_grad, k_grad, v_grad = q.grad, k.grad, v.grad + d_softmax_offset = None if is_training and config.softmax_type != "vanilla": d_softmax_offset = block.softmax_offset.grad if backend in ["UnfusedDotProductAttention"]: - if is_training: - return out, max_logit, (q.grad, k.grad, v.grad, d_softmax_offset) - else: - return out, max_logit, (None, None, None, d_softmax_offset) + return out, max_logit, (q_grad, k_grad, v_grad, d_softmax_offset) if backend in ["FusedAttention", "FlashAttention"]: if qkv_format == "thd" and pad_between_seqs: out_orig = torch.Tensor([]).to(device="cuda", dtype=dtype) @@ -1560,13 +1760,13 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: out_orig = torch.cat([out_orig, out[valid_range_q[0] : valid_range_q[1]]], dim=0) if is_training: q_grad_orig = torch.cat( - [q_grad_orig, q.grad[valid_range_q[0] : valid_range_q[1]]], dim=0 + [q_grad_orig, q_grad[valid_range_q[0] : valid_range_q[1]]], dim=0 ) k_grad_orig = torch.cat( - [k_grad_orig, k.grad[valid_range_kv[0] : valid_range_kv[1]]], dim=0 + [k_grad_orig, k_grad[valid_range_kv[0] : valid_range_kv[1]]], dim=0 ) v_grad_orig = torch.cat( - [v_grad_orig, v.grad[valid_range_kv[0] : valid_range_kv[1]]], dim=0 + [v_grad_orig, v_grad[valid_range_kv[0] : valid_range_kv[1]]], dim=0 ) if is_training: return ( @@ -1577,10 +1777,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: else: return out_orig, max_logit, (None, None, None, d_softmax_offset) else: - if is_training: - return out, max_logit, (q.grad, k.grad, v.grad, d_softmax_offset) - else: - return out, max_logit, (None, None, None, d_softmax_offset) + return out, max_logit, (q_grad, k_grad, v_grad, d_softmax_offset) model_configs_te_layer = { @@ -2196,8 +2393,20 @@ def get_model(dtype, config): } param_types_fp8_vs_f16 = [torch.float16, torch.bfloat16] -qkv_layout_fp8_vs_f16 = ["sbh3d", "bshd_bshd_bshd", "sbhd_sbhd_sbhd"] -qkv_format_fp8_vs_f16 = ["bshd", "sbhd"] +qkv_layout_fp8_vs_f16 = ["sbh3d", "bshd_bshd_bshd", "sbhd_sbhd_sbhd", "thd_thd_thd"] +qkv_format_fp8_vs_f16 = ["bshd", "sbhd", "thd"] + + +def _get_fp8_vs_f16_config(model, qkv_layout): + config = copy.copy(model_configs_fp8_vs_f16[model]) + # THD is variable-length, so it requires the corresponding padding-aware mask type. + if qkv_layout.startswith("thd"): + config.attn_mask_type = { + "no_mask": "padding", + "causal": "padding_causal", + "causal_bottom_right": "padding_causal_bottom_right", + }.get(config.attn_mask_type, config.attn_mask_type) + return config @pytest.mark.skipif(IS_HIP_EXTENSION, reason="FP8 Fused attention is not supported on ROCm") @@ -2223,7 +2432,7 @@ def test_mha_fp8_vs_f16( ): """Test MultiHeadAttention module in FP8""" os.environ["NVTE_FP8_DPA_BWD"] = "1" if fp8_dpa_bwd else "0" - config = model_configs_fp8_vs_f16[model] + config = _get_fp8_vs_f16_config(model, qkv_format) # Test backend availability if scaling_mode == "delayed": @@ -2383,19 +2592,32 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: if not is_training: mha = mha.eval() + def random_seqlens(max_seqlen): + if qkv_format != "thd": + return torch.randint( + 1, max_seqlen, [config.batch_size], dtype=torch.int32, device="cuda" + ) + # Reserve seven positions so total-token alignment only increases the final length. + return torch.cat( + ( + torch.randint( + 1, + max_seqlen, + [config.batch_size - 1], + dtype=torch.int32, + device="cuda", + ), + torch.randint(1, max_seqlen - 6, [1], dtype=torch.int32, device="cuda"), + ) + ) + if "padding" in config.attn_mask_type or qkv_format == "thd": if config.attn_type == "self": - seqlens_q = torch.randint( - 1, config.max_seqlen_q, [config.batch_size], dtype=torch.int32, device="cuda" - ) + seqlens_q = random_seqlens(config.max_seqlen_q) seqlens_kv = seqlens_q if config.attn_type == "cross": - seqlens_q = torch.randint( - 1, config.max_seqlen_q, [config.batch_size], dtype=torch.int32, device="cuda" - ) - seqlens_kv = torch.randint( - 1, config.max_seqlen_kv, [config.batch_size], dtype=torch.int32, device="cuda" - ) + seqlens_q = random_seqlens(config.max_seqlen_q) + seqlens_kv = random_seqlens(config.max_seqlen_kv) else: seqlens_q = torch.full( [config.batch_size], config.max_seqlen_q, dtype=torch.int32, device="cuda" @@ -2403,6 +2625,10 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: seqlens_kv = torch.full( [config.batch_size], config.max_seqlen_kv, dtype=torch.int32, device="cuda" ) + if qkv_format == "thd": + # FP8 Linear flattens THD input to [t, h*d], so align total tokens for cuBLAS. + seqlens_q[-1] += -seqlens_q.sum() % 8 + seqlens_kv[-1] += -seqlens_kv.sum() % 8 cu_seqlens_q = torch.zeros(config.batch_size + 1, dtype=torch.int32, device="cuda") cu_seqlens_kv = torch.zeros(config.batch_size + 1, dtype=torch.int32, device="cuda") cu_seqlens_q[1:] = torch.cumsum(seqlens_q, dim=0) @@ -2470,7 +2696,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: @pytest.mark.parametrize("scaling_mode", ["delayed", "current", "mxfp8"]) def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scaling_mode): """Test DotProductAttention module in FP8""" - config = model_configs_fp8_vs_f16[model] + config = _get_fp8_vs_f16_config(model, qkv_layout) # TODO(cyang): think of another way to verify dropout results # test cuDNN FP8 dropout @@ -2775,6 +3001,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: attn_mask_type=config.attn_mask_type, checkpoint_core_attention=False, core_attention_bias_type=config.attn_bias_type, + fp8_output=fp8_dpa, ) if is_training: out.backward(out_grad) diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 7911e9b334..0f6bb59da5 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -11,7 +11,6 @@ import subprocess import sys import threading -import time import pathlib import logging import copy @@ -35,7 +34,7 @@ from transformer_engine.pytorch.utils import get_torch_float8_e4m3_type _current_file = pathlib.Path(__file__).resolve() -sys.path.append(str(_current_file.parent.parent)) +sys.path = [str(_current_file.parent.parent)] + sys.path from utils import ModelConfig, get_available_attention_backends pytest_logging_level = logging.getLevelName(logging.root.level) @@ -175,8 +174,8 @@ def _kill(self) -> None: # One retry on pool-infrastructure failures (worker died / timed out / broken # pipe). Test-assertion failures from the worker carry the full per-rank - # traceback in resp["error"] and propagate without retry. Every retry leaves - # a [POOL-RETRY] line in stderr so pytest's capture surfaces + # traceback in resp["error"] and normally propagate without retry. Every retry + # leaves a [POOL-RETRY] line in stderr so pytest's capture surfaces # flake patterns in JUnit XML for offline analysis. _MAX_RETRIES = 1 @@ -186,17 +185,27 @@ def submit(self, kwargs: dict, timeout: float = POOL_SUBMIT_TIMEOUT_SEC) -> None try: return self._submit_once(kwargs, timeout) except AssertionError as e: - msg_head = str(e).splitlines()[0] + msg = str(e) + msg_head = msg.splitlines()[0] infrastructure_flake = ( "pool worker died" in msg_head or "timed out" in msg_head or "before request could be sent" in msg_head ) - if not infrastructure_flake or attempt == self._MAX_RETRIES: + # Heterogeneous CP cases can leave a retained worker in a state where + # FP8 THD emits NaNs even though the same case passes in a fresh worker. + # Retry only that signature once; a NaN from the fresh worker still fails. + fp8_thd_nan = ( + kwargs.get("dtype") == "fp8" + and kwargs.get("qkv_format") == "thd" + and "has nan values" in msg.lower() + ) + retryable = infrastructure_flake or fp8_thd_nan + if not retryable or attempt == self._MAX_RETRIES: if first_err is not None: sys.stderr.write( f"[POOL-RETRY-FAIL] world_size={self.world_size}: " - "both attempts died; first error was: " + "both attempts failed; first error was: " f"{str(first_err).splitlines()[0]!r}\n" ) sys.stderr.flush() @@ -204,7 +213,7 @@ def submit(self, kwargs: dict, timeout: float = POOL_SUBMIT_TIMEOUT_SEC) -> None first_err = e sys.stderr.write( f"[POOL-RETRY] world_size={self.world_size} attempt {attempt + 1} " - f"died: {msg_head!r}; respawning pool and retrying\n" + f"failed: {msg_head!r}; respawning pool and retrying\n" ) sys.stderr.flush() raise first_err # unreachable; loop either returns or raises @@ -306,7 +315,10 @@ def _submit(pool: PoolWorker, **kwargs) -> None: qkv_formats = ["sbhd", "thd"] -@pytest.mark.skipif(not FlashAttentionUtils.v2_plus, reason="Flash-attn 2.0+ is required.") +@pytest.mark.skipif( + not (FlashAttentionUtils.v2_plus or FlashAttentionUtils.v3_is_installed), + reason="Flash-attn v2 or v3 is required.", +) @pytest.mark.skipif(not IS_HIP_EXTENSION and get_device_compute_capability() < (8, 0), reason="CP tests require sm80+.") @pytest.mark.parametrize("dtype", dtypes) @pytest.mark.parametrize("model", model_configs_flash_attn.keys()) @@ -334,11 +346,6 @@ def test_cp_with_flash_attention(cp_pool, dtype, model, qkv_format, cp_comm_type if config.attn_bias_type != "no_bias" and cp_comm_type in ["all_gather", "a2a", "a2a+p2p"]: pytest.skip("No support for bias with cp_comm_type={all_gather, a2a, a2a+p2p}!") - if qkv_format == "thd" and cp_comm_type == "a2a+p2p": - pytest.skip( - "CP implementation with QKVO A2A+P2P (Hierarchical A2A) does not support THD format" - " yet!" - ) if ( qkv_format == "thd" and cp_comm_type == "all_gather" @@ -480,6 +487,8 @@ def test_cp_with_flash_attention(cp_pool, dtype, model, qkv_format, cp_comm_type "cp_4_3": ModelConfig( 2, 4096, 64, 64, attn_mask_type="causal", window_size=(128, 0), softmax_type="learnable" ), # GQA + "cp_5_0": ModelConfig(2, 1024, 16, 256, attn_mask_type="causal"), + "cp_5_1": ModelConfig(2, 1024, 16, 256, attn_mask_type="causal", window_size=(128, 0)), } @@ -498,6 +507,8 @@ def test_cp_with_flash_attention(cp_pool, dtype, model, qkv_format, cp_comm_type "cp_3_4", "cp_4_2", "cp_4_3", + "cp_5_0", + "cp_5_1", ] model_configs_fused_attn = {k: model_configs_fused_attn[k] for k in configs} dtypes = ["bf16", "fp8"] @@ -531,6 +542,23 @@ def test_cp_with_fused_attention( config.context_parallel = True config.cp_comm_type = cp_comm_type + if config.head_dim_qk == 256 and config.head_dim_v == 256: + # D=256 uses this generic CP runner, but only a subset of its axes is supported. + if get_device_compute_capability() not in ((10, 0), (10, 3)): + pytest.skip("D=256 CP fused attention is only enabled on Blackwell server GPUs.") + if dtype == "fp8": + pytest.skip("D=256 CP fused attention is covered for BF16/FP16 only.") + if cp_comm_type not in ["p2p", "all_gather"]: + pytest.skip("D=256 CP fused attention is covered for p2p and all_gather only.") + + required_cudnn_version = (9, 25, 0) if qkv_format == "thd" else (9, 23, 0) + required_cudnn_version_label = "9.25" if qkv_format == "thd" else "9.23" + if get_cudnn_version() < required_cudnn_version: + pytest.skip( + f"D=256 CP fused attention with {qkv_format.upper()} requires cuDNN" + f" {required_cudnn_version_label} or newer." + ) + num_gpus = 4 if cp_comm_type == "a2a+p2p" else 2 pool = cp_pool(num_gpus) @@ -550,8 +578,6 @@ def test_cp_with_fused_attention( if dtype != "fp8" and (fp8_mha or fp8_dpa): pytest.skip("dtype!=fp8 requires fp8_dpa=False and fp8_mha=False!") - if dtype == "fp8" and qkv_format == "thd": - pytest.skip("No support for FP8 attention with THD format!") if dtype == "fp8" and config.attn_bias_type != "no_bias": pytest.skip("No support for FP8 attention with bias!") @@ -560,10 +586,13 @@ def test_cp_with_fused_attention( if config.attn_bias_type != "no_bias" and cp_comm_type in ["all_gather", "a2a", "a2a+p2p"]: pytest.skip("No support for bias with cp_comm_type={all_gather, a2a, a2a+p2p}!") - if qkv_format == "thd" and cp_comm_type == "a2a+p2p": + # ROCm: upstream v2.19 lifted the a2a+p2p (Hierarchical A2A) THD skip after + # implementing cuDNN support for it. ROCm's CK fused-attn backend has no + # equivalent path, so keep the skip gated to ROCm until CK gains support. + if IS_HIP_EXTENSION and qkv_format == "thd" and cp_comm_type == "a2a+p2p": pytest.skip( "CP implementation with QKVO A2A+P2P (Hierarchical A2A) does not support THD format" - " yet!" + " yet on ROCm!" ) # ROCm: upstream v2.18 narrowed the fused-path THD skip from @@ -579,6 +608,7 @@ def test_cp_with_fused_attention( " seqused_k path)." ) + if (config.window_size[0] != -1 or config.window_size[1] not in [-1, 0]) and cp_comm_type in [ "p2p", "a2a+p2p", @@ -612,6 +642,8 @@ def test_cp_with_fused_attention( pytest.skip("scaling_mode=delayed requires f16_O=False!") if scaling_mode == "mxfp8" and not f16_O: pytest.skip("scaling_mode=mxfp8 requires f16_O=True!") + if scaling_mode == "mxfp8" and qkv_format == "thd": + pytest.skip("MXFP8 quantization does not support THD format!") if scaling_mode == "mxfp8" and fp8_mha: pytest.skip("No support for scaling_mode=mxfp8 with fp8_mha=True!") diff --git a/tests/pytorch/attention/test_cp_utils.py b/tests/pytorch/attention/test_cp_utils.py index c3a423cef5..f63f659aac 100644 --- a/tests/pytorch/attention/test_cp_utils.py +++ b/tests/pytorch/attention/test_cp_utils.py @@ -8,10 +8,12 @@ import torch import unittest from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import ( + _zero_thd_padding, get_batch_on_this_cp_rank, pad_thd_sequences_for_cp, generate_positional_ids_for_cp, ) +from transformer_engine.pytorch.attention.dot_product_attention.utils import get_thd_padding_mask try: import transformer_engine_torch as tex @@ -809,6 +811,84 @@ def _legacy_valid_copy(out, inp, cu_seqlens_padded, cu_seqlens): out[s : s + sz].copy_(inp[s : s + sz]) +@unittest.skipIf(not torch.cuda.is_available(), "THD padding-mask tests require CUDA") +class TestTHDPaddingMask(unittest.TestCase): + @staticmethod + def _reference_mask(cu_seqlens, cu_seqlens_padded): + mask = torch.ones(cu_seqlens_padded[-1].item(), dtype=torch.bool) + for batch_idx in range(cu_seqlens.numel() - 1): + start = cu_seqlens_padded[batch_idx].item() + length = (cu_seqlens[batch_idx + 1] - cu_seqlens[batch_idx]).item() + mask[start : start + length] = False + return mask.cuda() + + def test_matches_reference_across_batch_sizes_and_q_kv_layouts(self): + for batch_size in (2, 8, 32, 128): + with self.subTest(batch_size=batch_size): + sequence = torch.arange(batch_size, dtype=torch.int32) + layouts = ( + (17 + sequence % 7, torch.full_like(sequence, 32)), + (9 + sequence * 3 % 11, torch.full_like(sequence, 24)), + ) + for seqlens, padded_seqlens in layouts: + cu_seqlens = torch.cat((torch.zeros(1, dtype=torch.int32), seqlens.cumsum(0))) + cu_seqlens_padded = torch.cat( + (torch.zeros(1, dtype=torch.int32), padded_seqlens.cumsum(0)) + ) + expected = self._reference_mask(cu_seqlens, cu_seqlens_padded) + cu_seqlens = cu_seqlens.cuda() + cu_seqlens_padded = cu_seqlens_padded.cuda() + actual = get_thd_padding_mask(expected.numel(), cu_seqlens, cu_seqlens_padded) + self.assertTrue(torch.equal(actual, expected)) + + def test_zeroes_padding_without_changing_valid_rows(self): + cu_seqlens = torch.tensor([0, 3, 8], dtype=torch.int32) + cu_seqlens_padded = torch.tensor([0, 4, 12], dtype=torch.int32) + padding_mask = self._reference_mask(cu_seqlens, cu_seqlens_padded) + cu_seqlens = cu_seqlens.cuda() + cu_seqlens_padded = cu_seqlens_padded.cuda() + + tensors = tuple( + torch.arange(1, 25, dtype=torch.float32, device="cuda").view(12, 2) + offset + for offset in (0, 100) + ) + originals = tuple(tensor.clone() for tensor in tensors) + + _zero_thd_padding(tensors, cu_seqlens, cu_seqlens_padded) + + for tensor, original in zip(tensors, originals): + self.assertTrue(torch.equal(tensor[~padding_mask], original[~padding_mask])) + self.assertEqual(torch.count_nonzero(tensor[padding_mask]).item(), 0) + + def test_zero_padding_is_cuda_graph_safe(self): + cu_seqlens = torch.tensor([0, 3, 8], dtype=torch.int32, device="cuda") + cu_seqlens_padded = torch.tensor([0, 4, 12], dtype=torch.int32, device="cuda") + tensors = tuple(torch.ones((12, 2), dtype=torch.float32, device="cuda") for _ in range(2)) + + side_stream = torch.cuda.Stream() + side_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side_stream): + _zero_thd_padding(tensors, cu_seqlens, cu_seqlens_padded) + torch.cuda.current_stream().wait_stream(side_stream) + torch.cuda.synchronize() + + for tensor in tensors: + tensor.fill_(1) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _zero_thd_padding(tensors, cu_seqlens, cu_seqlens_padded) + + for tensor in tensors: + tensor.fill_(1) + graph.replay() + torch.cuda.synchronize() + + padding_mask = get_thd_padding_mask(12, cu_seqlens, cu_seqlens_padded) + for tensor in tensors: + self.assertEqual(torch.count_nonzero(tensor[padding_mask]).item(), 0) + self.assertTrue(torch.all(tensor[~padding_mask] == 1)) + + @unittest.skipIf( not torch.cuda.is_available() or tex is None, "THD kernel tests require CUDA and transformer_engine_torch", @@ -881,6 +961,22 @@ def test_thd_read_half_tensor_reads_each_sequence_half(self): torch.equal(kv_second, torch.stack([expected_second, expected_second + 128])) ) + def test_thd_grad_correction_copies_byte_half_and_zeros_inactive_half(self): + cu_seqlens = torch.tensor([0, 8, 20], dtype=torch.int32, device="cuda") + grad_per_step = torch.arange(10 * 2 * 8, dtype=torch.uint8, device="cuda").view(10, 2, 8) + first_half_rows = torch.tensor([0, 1, 2, 3, 8, 9, 10, 11, 12, 13], device="cuda") + second_half_rows = torch.tensor([4, 5, 6, 7, 14, 15, 16, 17, 18, 19], device="cuda") + + grad = torch.full((20, 2, 8), 255, dtype=torch.uint8, device="cuda") + tex.thd_grad_correction(grad, grad_per_step, cu_seqlens, "copy", "zero") + self.assertTrue(torch.equal(grad[first_half_rows], grad_per_step)) + self.assertEqual(torch.count_nonzero(grad[second_half_rows]).item(), 0) + + grad.fill_(255) + tex.thd_grad_correction(grad, grad_per_step, cu_seqlens, "zero", "copy") + self.assertEqual(torch.count_nonzero(grad[first_half_rows]).item(), 0) + self.assertTrue(torch.equal(grad[second_half_rows], grad_per_step)) + def test_thd_read_second_half_lse_handles_packed_and_batch_major_lse(self): cu_seqlens = torch.tensor([0, 8, 16], dtype=torch.int32, device="cuda") lse = torch.arange(2 * 2 * 8, dtype=torch.float32, device="cuda").view(2, 2, 8) diff --git a/tests/pytorch/attention/test_fused_mla_q_uproj.py b/tests/pytorch/attention/test_fused_mla_q_uproj.py new file mode 100644 index 0000000000..2c061607fc --- /dev/null +++ b/tests/pytorch/attention/test_fused_mla_q_uproj.py @@ -0,0 +1,137 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Unit tests for FusedMLAQUpProjRopeQuant. + +Run: + pytest tests/pytorch/attention/test_fused_mla_q_uproj.py -v +""" + +import pytest +import torch + +import transformer_engine.pytorch # registers transformer_engine_torch +import transformer_engine_torch as tex +from transformer_engine.pytorch.attention import FusedMLAQUpProjRopeQuant +from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor + +# DSv3 671B MLA dims +NUM_HEADS = 128 +HEAD_DIM_NOPE = 128 +HEAD_DIM_ROPE = 64 +HEAD_DIM = HEAD_DIM_NOPE + HEAD_DIM_ROPE # 192 +Q_LORA_RANK = 1536 +PROJ_DIM = NUM_HEADS * HEAD_DIM # 24576 + +SEED = 42 + +fused_supported, reason_not_supported = ( + (True, "") + if FusedMLAQUpProjRopeQuant.is_supported() + else ( + False, + ( + "FusedMLAQUpProjRopeQuant.is_supported() returned False " + "(SM100+, cudnn-frontend >= 1.27.0, and NVTE_FUSED_MLA_Q_UPROJ=1 required)" + ), + ) +) + + +def _dequantize_fused_output(query: MXFP8Tensor, s: int, b: int) -> torch.Tensor: + """Dequantize the rowwise fused output to bf16 [s, b, nh, head_dim]. + + TE's C++ dequantize kernel requires 2D layout, so reshape before calling dequantize(). + """ + tokens = s * b + q_2d = MXFP8Tensor( + shape=(tokens, PROJ_DIM), + dtype=torch.bfloat16, + rowwise_data=query._rowwise_data.view(tokens, PROJ_DIM), + rowwise_scale_inv=query._rowwise_scale_inv.view(tokens, PROJ_DIM // 32), + columnwise_data=None, + columnwise_scale_inv=None, + quantizer=query._quantizer, + requires_grad=False, + fp8_dtype=query._fp8_dtype, + with_gemm_swizzled_scales=False, + ) + return q_2d.dequantize().to(torch.bfloat16).view(s, b, NUM_HEADS, HEAD_DIM) + + +def _reference_q_uproj( + x: torch.Tensor, + w_mxfp8: MXFP8Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + s: int, + b: int, +) -> torch.Tensor: + """Unfused bf16 reference: dequantize-then-GEMM + RoPE. Returns [s, b, nh, head_dim] bf16.""" + x_dq = ( + MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False)(x) + .dequantize() + .to(torch.bfloat16) + ) + w_dq = w_mxfp8.dequantize().to(torch.bfloat16) + out = (x_dq @ w_dq.t()).view(s, b, NUM_HEADS, HEAD_DIM) + + q_nope = out[..., :HEAD_DIM_NOPE] + q_rope = out[..., HEAD_DIM_NOPE:] + cos_ = cos[:, None, None, :].to(q_rope.dtype) + sin_ = sin[:, None, None, :].to(q_rope.dtype) + half = HEAD_DIM_ROPE // 2 + x1, x2 = q_rope[..., 0::2], q_rope[..., 1::2] + q_rope_out = torch.cat( + [ + x1 * cos_[..., :half] - x2 * sin_[..., :half], + x2 * cos_[..., half:] + x1 * sin_[..., half:], + ], + dim=-1, + ) + return torch.cat([q_nope, q_rope_out], dim=-1) + + +def _build_rope_tables(tokens: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: + inv_freq = 1.0 / ( + 10000 + ** (torch.arange(0, HEAD_DIM_ROPE, 2, dtype=torch.float32, device=device) / HEAD_DIM_ROPE) + ) + freqs = torch.cat( + [torch.outer(torch.arange(tokens, device=device, dtype=torch.float32), inv_freq)] * 2, + dim=-1, + ) + return freqs.cos().to(torch.bfloat16), freqs.sin().to(torch.bfloat16) + + +@pytest.mark.skipif(not fused_supported, reason=reason_not_supported) +@pytest.mark.parametrize("tokens", [256]) +def test_fused_mla_q_uproj(tokens: int) -> None: + """Forward numerics and x_saved properties for FusedMLAQUpProjRopeQuant.run(). + + Full forward+backward autograd testing (via _FusedMLAQUpProjFunction) lives in + Megatron-Core. + """ + s, b = tokens, 1 + device = torch.device("cuda") + torch.manual_seed(SEED) + torch.cuda.manual_seed(SEED) + + x = torch.randn(tokens, Q_LORA_RANK, dtype=torch.bfloat16, device=device) + w = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False)( + torch.randn(PROJ_DIM, Q_LORA_RANK, dtype=torch.bfloat16, device=device) + ) + cos, sin = _build_rope_tables(tokens, device) + + query, x_saved = FusedMLAQUpProjRopeQuant.run(x, w, cos, sin, s, b) + + # Forward numerics: FP8 GEMM + output quantize introduce ~10% relative error. + fused_dq = _dequantize_fused_output(query, s, b) + ref_dq = _reference_q_uproj(x, w, cos, sin, s, b) + torch.testing.assert_close(fused_dq, ref_dq, atol=0.5, rtol=0.1) + + # x_saved: must be MXFP8 with only columnwise data retained for wgrad. + assert isinstance(x_saved, MXFP8Tensor) + assert x_saved._columnwise_data is not None, "x_saved must retain columnwise data for wgrad" + assert x_saved._rowwise_data is None, "x_saved rowwise data should be dropped after forward" diff --git a/tests/pytorch/attention/test_kv_cache.py b/tests/pytorch/attention/test_kv_cache.py index cc24056d97..1d6a8cbf34 100644 --- a/tests/pytorch/attention/test_kv_cache.py +++ b/tests/pytorch/attention/test_kv_cache.py @@ -38,7 +38,8 @@ ) _current_file = pathlib.Path(__file__).resolve() -sys.path.append(str(_current_file.parent.parent)) +# Prepend so installed packages with a top-level utils module cannot shadow the test helpers. +sys.path = [str(_current_file.parent.parent)] + sys.path from utils import ( ModelConfig, reset_rng_states, diff --git a/tests/pytorch/debug/test_sanity.py b/tests/pytorch/debug/test_sanity.py index 2bc4b35590..c3cb5237af 100644 --- a/tests/pytorch/debug/test_sanity.py +++ b/tests/pytorch/debug/test_sanity.py @@ -7,6 +7,20 @@ import nvdlfw_inspect.api as debug_api import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine.debug.pytorch.debug_quantization import ( + DebugQuantizedTensor, + DebugQuantizer, + HIGH_PRECISION, + STANDARD_QUANTIZE, +) +from transformer_engine.pytorch.module._common import ( + set_quantizer_usage_for_wgrad_all_gather, +) +from transformer_engine.pytorch.quantized_tensor import Quantizer +from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockQuantizer +from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer +from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer from test_numerics import create_config_file @@ -58,6 +72,113 @@ """, } + +def _make_debug_quantizer_for_usage_test(parent_quantizer): + """Construct a DebugQuantizer without initializing the debug API.""" + quantizer = object.__new__(DebugQuantizer) + Quantizer.__init__(quantizer, rowwise=True, columnwise=True) + quantizer.parent_quantizer = parent_quantizer + quantizer.output_tensor = False + if parent_quantizer is None: + quantizer.rowwise_tensor_plan = HIGH_PRECISION + quantizer.columnwise_tensor_plan = HIGH_PRECISION + else: + quantizer.rowwise_tensor_plan = STANDARD_QUANTIZE + quantizer.columnwise_tensor_plan = STANDARD_QUANTIZE + return quantizer + + +def test_wgrad_all_gather_usage_handles_debug_quantizer_without_parent(): + """High-precision debug mode has no parent quantizer to unwrap.""" + quantizer = _make_debug_quantizer_for_usage_test(None) + + set_quantizer_usage_for_wgrad_all_gather(quantizer) + + assert quantizer.rowwise_usage is False + assert quantizer.columnwise_usage is True + + +def test_wgrad_all_gather_usage_updates_debug_wrapper_and_blockwise_parent(): + """Usage changes must remain synchronized across the debug wrapper and parent.""" + parent_quantizer = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=1, + ) + quantizer = _make_debug_quantizer_for_usage_test(parent_quantizer) + + set_quantizer_usage_for_wgrad_all_gather(quantizer) + + assert quantizer.rowwise_usage is False + assert quantizer.columnwise_usage is True + assert parent_quantizer.rowwise_usage is False + assert parent_quantizer.columnwise_usage is True + + +def test_wgrad_all_gather_usage_detects_hybrid_through_debug_wrapper(): + """Hybrid classification may inspect the parent, but usage mutates the wrapper.""" + parent_quantizer = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + quantizer = _make_debug_quantizer_for_usage_test(parent_quantizer) + + set_quantizer_usage_for_wgrad_all_gather(quantizer) + + assert quantizer.rowwise_usage is False + assert quantizer.columnwise_usage is True + assert parent_quantizer.rowwise_usage is False + assert parent_quantizer.columnwise_usage is True + + +def test_debug_quantized_tensor_routes_usage_to_distinct_representations(): + """A usage request should only reach the child for that GEMM direction.""" + rowwise = IdentityQuantizer()(torch.ones(2, 2)) + columnwise = IdentityQuantizer()(torch.ones(2, 2)) + rowwise_calls = [] + columnwise_calls = [] + + def record_rowwise(rowwise_usage=None, columnwise_usage=None): + rowwise_calls.append((rowwise_usage, columnwise_usage)) + + def record_columnwise(rowwise_usage=None, columnwise_usage=None): + columnwise_calls.append((rowwise_usage, columnwise_usage)) + + rowwise.update_usage = record_rowwise + columnwise.update_usage = record_columnwise + tensor = DebugQuantizedTensor(rowwise, columnwise, quantizer=None) + + tensor.update_usage(columnwise_usage=True) + assert rowwise_calls == [] + assert columnwise_calls == [(None, True)] + + tensor.update_usage(rowwise_usage=True) + assert rowwise_calls == [(True, None)] + assert columnwise_calls == [(None, True)] + + tensor.update_usage(rowwise_usage=False) + assert tensor.rowwise_gemm_tensor is None + with pytest.raises(RuntimeError, match="Cannot recreate rowwise tensor"): + tensor.update_usage(rowwise_usage=True, columnwise_usage=False) + assert tensor.columnwise_gemm_tensor is columnwise + + +def test_debug_quantized_tensor_updates_shared_representation_once(): + """A shared child receives one combined usage update.""" + shared = IdentityQuantizer()(torch.ones(2, 2)) + calls = [] + + def record_usage(rowwise_usage=None, columnwise_usage=None): + calls.append((rowwise_usage, columnwise_usage)) + + shared.update_usage = record_usage + tensor = DebugQuantizedTensor(shared, shared, quantizer=None) + + tensor.update_usage(rowwise_usage=True, columnwise_usage=False) + assert calls == [(True, False)] + + # Configs that require FP8 to be enabled fp8_required_configs = {"log_fp8"} diff --git a/tests/pytorch/distributed/fsdp2_tests/conftest.py b/tests/pytorch/distributed/fsdp2_tests/conftest.py index 07e264f54d..12f98158b4 100644 --- a/tests/pytorch/distributed/fsdp2_tests/conftest.py +++ b/tests/pytorch/distributed/fsdp2_tests/conftest.py @@ -14,6 +14,7 @@ import torch import torch.distributed as dist from transformer_engine.pytorch import fp8 +from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported # Ensure the correct CUDA device is active before _parametrize_recipes() # runs at collection time, since the session-scoped dist_init fixture @@ -45,6 +46,13 @@ def _check_nvfp4_support(): ("NVFP4BlockScaling", _check_nvfp4_support), ] +_HYBRID_RECIPE_CONFIGS = [ + ("HybridFP8CurrentScaling", fp8.check_fp8_support), + ("HybridMXFP8", fp8.check_mxfp8_support), + ("HybridFloat8BlockScaling", fp8.check_fp8_block_scaling_support), + ("HybridMixed_MXFP8_FP8", fp8.check_mxfp8_support), +] + def _parametrize_recipes(): params = [] @@ -56,6 +64,26 @@ def _parametrize_recipes(): return params +def _parametrize_hybrid_recipes(): + params = [] + for name, check_fn in _HYBRID_RECIPE_CONFIGS: + supported, reason = check_fn() + marks = [pytest.mark.skipif(not supported, reason=reason)] + if name == "HybridFP8CurrentScaling" and not is_non_tn_fp8_gemm_supported(): + marks.append( + pytest.mark.xfail( + raises=NotImplementedError, + strict=True, + reason=( + "Hopper does not yet support columnwise-only per-tensor FP8 " + "quantization; tracked by NVIDIA/TransformerEngine#3158" + ), + ) + ) + params.append(pytest.param(name, id=name, marks=marks)) + return params + + # ── Session / per-test fixtures ────────────────────────────────────── @pytest.fixture(scope="session", autouse=True) def dist_init(): @@ -88,3 +116,8 @@ def _cleanup(): @pytest.fixture(params=_parametrize_recipes()) def recipe_name(request): return request.param + + +@pytest.fixture(params=_parametrize_hybrid_recipes()) +def hybrid_recipe_name(request): + return request.param diff --git a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py index 178ce62375..81cc8731f5 100644 --- a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py +++ b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py @@ -4,14 +4,69 @@ """Shared utility functions for FSDP2 distributed tests.""" +import sys +from pathlib import Path + +# FSDP2 files are collected directly by subprocess pytest invocations, which +# put this directory (rather than tests/pytorch) on sys.path. +_TEST_ROOT = str(Path(__file__).resolve().parents[2]) +if _TEST_ROOT not in sys.path: + sys.path.append(_TEST_ROOT) + import transformer_engine.common.recipe from transformer_engine.pytorch import QuantizedTensor +from hybrid_quantization_utils import ( + hybrid_float8_block_qfactory, + hybrid_fp8_current_identity_qfactory, + hybrid_fp8_current_qfactory, + hybrid_mixed_mxfp8_fp8_qfactory, + hybrid_mxfp8_qfactory, + identity_qfactory, +) + def get_recipe_from_string(recipe): return getattr(transformer_engine.common.recipe, recipe)() +# CustomRecipe has dynamic TE extra-state handling. Once FP8 state is +# initialized, TE's get_extra_state() pickles the recipe on save, so +# checkpoint-test qfactories must be module-level and picklable. On load, +# payloads without delayed-scaling state are identified and ignored without +# unpickling. See ``run_fsdp2_fused_adam.py::test_hybrid_dcp_output_parity``. +_HYBRID_QFACTORIES = { + "HybridFP8CurrentScaling": hybrid_fp8_current_qfactory, + "HybridMXFP8": hybrid_mxfp8_qfactory, + "HybridFloat8BlockScaling": hybrid_float8_block_qfactory, + "HybridMixed_MXFP8_FP8": hybrid_mixed_mxfp8_fp8_qfactory, + "HybridFP8CurrentScalingIdentity": hybrid_fp8_current_identity_qfactory, + "Identity": identity_qfactory, +} + + +def get_hybrid_recipe_from_string(recipe): + """Build a CustomRecipe wrapping a module-level (picklable) hybrid qfactory. + + Each hybrid qfactory composes one or two role-aware base factories from + ``quantizer_factories`` per direction; per-role behavior is delegated + to the base factory and the hybrid layer only decides the direction pairing. + + Supported values: + "HybridFP8CurrentScaling" — FP8 current for both directions + "HybridMXFP8" — MXFP8 for both directions + "HybridFloat8BlockScaling" — Float8 block scaling for both directions + "HybridMixed_MXFP8_FP8" — MXFP8 rowwise + FP8 current columnwise + "HybridFP8CurrentScalingIdentity" — FP8 current forward + Identity backward + "Identity" — high-precision passthrough for every slot + """ + if recipe not in _HYBRID_QFACTORIES: + raise ValueError( + f"Unknown hybrid recipe '{recipe}'. Supported: {sorted(_HYBRID_QFACTORIES.keys())}" + ) + return transformer_engine.common.recipe.CustomRecipe(qfactory=_HYBRID_QFACTORIES[recipe]) + + def save_custom_attrs(module): custom_attrs = {} for name, param in module.named_parameters(): diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index c27eb21d4d..db14af57b5 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -47,7 +47,12 @@ from torch.distributed.tensor import DTensor import transformer_engine.pytorch as te -from transformer_engine.pytorch import QuantizedTensor +from transformer_engine.pytorch import HybridQuantizedTensor, QuantizedTensor +from transformer_engine.pytorch.tensor import ( + Float8BlockwiseQTensorStorage, + Float8TensorStorage, + MXFP8TensorStorage, +) import transformer_engine.common.recipe # Executed as a script, so sibling imports rely on the interpreter putting this file's @@ -55,7 +60,6 @@ sys.path.append(os.path.dirname(os.path.realpath(__file__))) from fsdp2_utils import get_recipe_from_string, save_custom_attrs, restore_custom_attrs - HIDDEN_SIZE = 256 FFN_HIDDEN_SIZE = 1024 NUM_ATTENTION_HEADS = 8 @@ -119,7 +123,7 @@ def _build_model( return model -def _shard_model(model, world_size): +def _shard_model(model, world_size, reshard_after_forward=None): """Apply FSDP2 sharding with save/restore custom attrs. If the model was created on the meta device (e.g. for FP8 init), @@ -128,12 +132,24 @@ def _shard_model(model, world_size): restore_custom_attrs is called last so it applies to the final parameter objects. For meta-device models, reset_parameters() replaces params via module_setattr (base.py:1336-1339), so attrs must be restored afterward. + + Parameters + ---------- + reshard_after_forward : bool, optional + Passed through to ``fully_shard``. ``None`` (default) keeps FSDP2's + own default: ``True`` for child modules, ``False`` for the root. + ``False`` on child modules keeps the full-precision gathered weight + alive through backward, exercising the iter-2+ buffer-reuse path + inside the same forward/backward rather than across training steps. """ has_meta_params = any(p.is_meta for p in model.parameters()) custom_attrs = save_custom_attrs(model) mesh = DeviceMesh("cuda", list(range(world_size))) + shard_kwargs = {"mesh": mesh} + if reshard_after_forward is not None: + shard_kwargs["reshard_after_forward"] = reshard_after_forward for child in model.children(): - fully_shard(child, mesh=mesh) + fully_shard(child, **shard_kwargs) fully_shard(model, mesh=mesh) if has_meta_params: for module in model.modules(): @@ -154,6 +170,185 @@ def _get_dist_info(): return world_size, device +def _collective_assert(condition, message): + """Raise on every rank only after all ranks report structural status.""" + failed = torch.tensor([not condition], dtype=torch.uint8, device="cuda") + if not condition: + print(f"[rank {dist.get_rank()}] {message}", flush=True) + dist.all_reduce(failed, dist.ReduceOp.MAX) + assert not bool(failed.item()), f"{message}: failed on at least one rank" + + +def _collective_assert_same_shape(tensor, message): + """Check tensor ranks/shapes before entering a shape-sensitive collective.""" + # Quantized parameter buffers in these tests are at most 2-D, but leave + # room for future formats without introducing object collectives. + layout = torch.full((9,), -1, dtype=torch.int64, device=tensor.device) + _collective_assert(tensor.ndim < layout.numel(), f"{message}: ndim={tensor.ndim}") + layout[0] = tensor.ndim + layout[1 : tensor.ndim + 1] = torch.tensor( + tensor.shape, dtype=torch.int64, device=tensor.device + ) + min_layout = layout.clone() + max_layout = layout.clone() + dist.all_reduce(min_layout, dist.ReduceOp.MIN) + dist.all_reduce(max_layout, dist.ReduceOp.MAX) + _collective_assert( + torch.equal(min_layout, max_layout), + f"{message}: tensor shape differs across ranks", + ) + + +def _record_exact(errors, actual, expected, label): + """Record an exact mismatch without interrupting later distributed work.""" + if actual is None or expected is None: + if actual is not expected: + errors.append(f"{label}: one tensor is None") + return + try: + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + except (AssertionError, TypeError) as exc: + errors.append(f"{label}: {exc}") + + +def _raise_collective_errors(errors, context): + """Aggregate delayed exact failures after every rank finishes collectives.""" + if errors: + print(f"[rank {dist.get_rank()}] {context}:\n" + "\n".join(errors), flush=True) + failed = torch.tensor([bool(errors)], dtype=torch.uint8, device="cuda") + dist.all_reduce(failed, dist.ReduceOp.MAX) + assert not bool(failed.item()), f"{context}: exact check failed on at least one rank" + + +def _check_hybrid_direction_buffers( + local_sub, + full_sub, + expected_type, + *, + direction, + param_name, + world_size, + errors, +): + """Compare one direction's gathered raw data, scale buffers, and metadata.""" + _collective_assert( + isinstance(local_sub, expected_type), + f"{param_name}: {direction} local storage is {type(local_sub).__name__}, " + f"expected {expected_type.__name__}", + ) + _collective_assert( + isinstance(full_sub, expected_type), + f"{param_name}: {direction} full storage is {type(full_sub).__name__}, " + f"expected {expected_type.__name__}", + ) + + local_buffers, local_meta = local_sub.fsdp_extract_buffers() + full_buffers, full_meta = full_sub.fsdp_extract_buffers() + local_count = torch.tensor([len(local_buffers)], dtype=torch.int64, device="cuda") + min_count = local_count.clone() + max_count = local_count.clone() + dist.all_reduce(min_count, dist.ReduceOp.MIN) + dist.all_reduce(max_count, dist.ReduceOp.MAX) + _collective_assert( + min_count.item() == max_count.item() == len(full_buffers), + f"{param_name}: {direction} buffer count mismatch: local={len(local_buffers)}, " + f"full={len(full_buffers)}, rank range=({min_count.item()}, {max_count.item()})", + ) + + expected_buffers = [] + for buffer in local_buffers: + _collective_assert( + buffer is not None, + f"{param_name}: {direction} unexpectedly exposed a None FSDP buffer", + ) + _collective_assert_same_shape(buffer, f"{param_name}: {direction} FSDP buffer") + gathered = [torch.zeros_like(buffer) for _ in range(world_size)] + dist.all_gather(gathered, buffer) + expected_buffers.append(torch.cat(gathered, dim=0)) + + for key in ("field_names", "direction"): + if key in local_meta or key in full_meta: + if local_meta.get(key) != full_meta.get(key): + errors.append( + f"{param_name}: {direction} metadata {key!r} differs: " + f"{full_meta.get(key)!r} != {local_meta.get(key)!r}" + ) + for index, (actual, expected) in enumerate(zip(full_buffers, expected_buffers)): + field_names = local_meta.get("field_names", ()) + field = field_names[index] if index < len(field_names) else f"buffer[{index}]" + _record_exact(errors, actual, expected, f"{param_name}: {direction} {field}") + + # Per-tensor FP8 scale is metadata rather than an FSDP buffer. It must be + # identical on every shard and preserved on the reconstructed full tensor. + if isinstance(local_sub, Float8TensorStorage): + local_scale_inv = getattr(local_sub, "_scale_inv", None) + _collective_assert( + isinstance(local_scale_inv, torch.Tensor), + f"{param_name}: {direction} Float8 storage has no tensor _scale_inv", + ) + local_scale = local_scale_inv.detach().clone() + _collective_assert_same_shape(local_scale, f"{param_name}: {direction} Float8 scale") + gathered_scales = [torch.zeros_like(local_scale) for _ in range(world_size)] + dist.all_gather(gathered_scales, local_scale) + for rank, scale in enumerate(gathered_scales[1:], start=1): + _record_exact( + errors, + scale, + gathered_scales[0], + f"{param_name}: {direction} scale rank {rank}", + ) + _record_exact( + errors, + full_sub._scale_inv, + gathered_scales[0], + f"{param_name}: {direction} reconstructed scale", + ) + + +def _manual_reconstruct_hybrid(local, *, param_name, world_size): + """Run the Hybrid FSDP pre/post protocol on manually gathered raw buffers.""" + _collective_assert( + isinstance(local, HybridQuantizedTensor), + f"{param_name}: local shard is {type(local).__name__}, expected HybridQuantizedTensor", + ) + sharded_tensors, metadata = local.fsdp_pre_all_gather( + mesh=None, + orig_size=local.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + local_count = torch.tensor([len(sharded_tensors)], dtype=torch.int64, device=local.device) + min_count = local_count.clone() + max_count = local_count.clone() + dist.all_reduce(min_count, dist.ReduceOp.MIN) + dist.all_reduce(max_count, dist.ReduceOp.MAX) + _collective_assert( + min_count.item() == max_count.item() and min_count.item() > 0, + f"{param_name}: Hybrid FSDP buffer count range is ({min_count.item()}, {max_count.item()})", + ) + + gathered_outputs = [] + for index, shard in enumerate(sharded_tensors): + _collective_assert( + shard is not None, + f"{param_name}: Hybrid FSDP buffer {index} is None", + ) + _collective_assert_same_shape(shard, f"{param_name}: Hybrid FSDP buffer {index}") + gathered = [torch.zeros_like(shard) for _ in range(world_size)] + dist.all_gather(gathered, shard) + gathered_outputs.append(torch.cat(gathered, dim=0)) + + reconstructed, _ = local.fsdp_post_all_gather( + tuple(gathered_outputs), metadata, local.dtype, out=None + ) + _collective_assert( + isinstance(reconstructed, HybridQuantizedTensor), + f"{param_name}: Hybrid FSDP reconstruction returned {type(reconstructed).__name__}", + ) + return reconstructed + + def test_fused_adam_fp8_master_weights(recipe_name): """FusedAdam with master_weights + FSDP2 + quantized_model_init (FP8 params). @@ -1127,6 +1322,1047 @@ def test_dcp_resharding_load(recipe_name): os.remove(ref_output_path) +# --------------------------------------------------------------------------- +# Hybrid quantization + FSDP2 tests +# --------------------------------------------------------------------------- + + +def _build_hybrid_model(hybrid_recipe, use_meta_device=True): + """Build a model with quantized_model_init using a hybrid CustomRecipe.""" + kwargs = dict( + fuse_qkv_params=True, + params_dtype=torch.bfloat16, + hidden_dropout=0.0, + attention_dropout=0.0, + ) + if use_meta_device: + kwargs["device"] = "meta" + with te.quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = torch.nn.Sequential( + *[ + te.TransformerLayer( + HIDDEN_SIZE, + FFN_HIDDEN_SIZE, + NUM_ATTENTION_HEADS, + **kwargs, + ) + for _ in range(NUM_LAYERS) + ] + ) + return model + + +def test_fused_adam_hybrid_master_weights(hybrid_recipe_name): + """FusedAdam + master_weights + FSDP2 + hybrid quantized_model_init. + + Verifies: + - Params are DTensors wrapping HybridQuantizedTensor local shards + - Training loop completes without error + - Optimizer states are FP32 + - Loss decreases over training steps + """ + from transformer_engine.pytorch import HybridQuantizedTensor + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size, device = _get_dist_info() + + model = _build_hybrid_model(hybrid_recipe) + model = _shard_model(model, world_size) + + hybrid_count = sum( + 1 + for _, p in model.named_parameters() + if isinstance(p, DTensor) and isinstance(p._local_tensor, HybridQuantizedTensor) + ) + assert hybrid_count > 0, "No HybridQuantizedTensor local tensors after sharding" + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + losses = [] + for step in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + loss = F.mse_loss(output, target) + losses.append(loss.item()) + loss.backward() + optimizer.step() + + for param in model.parameters(): + state = optimizer.state[param] + assert state["exp_avg"].dtype == torch.float32 + assert state["exp_avg_sq"].dtype == torch.float32 + if "master_param" in state: + assert state["master_param"].dtype == torch.float32 + + # Strictly monotonic decrease + assert all( + losses[i + 1] < losses[i] for i in range(len(losses) - 1) + ), f"Loss not strictly decreasing each step: {losses}" + + +def test_fused_adam_hybrid_reshard_variants(hybrid_recipe_name): + """Hybrid FusedAdam training must be numerically invariant to FSDP2's + ``reshard_after_forward`` schedule. + + ``reshard_after_forward`` only changes *when* the gathered weight is + materialized/freed, not the math: ``True`` (FSDP2's child-module default) + drops the gathered weight after forward and re-gathers it in backward -- + invoking ``fsdp_post_all_gather(out=...)`` twice per step -- while ``False`` + keeps the gathered copy alive through backward (one gather per step). The + gathered quantized bytes are identical either way, so both schedules must + produce bitwise-identical outputs, losses, input/parameter gradients, + master parameters, moments, and step counters. + + Strictly stronger than "loss decreased": it locks in that the hybrid + all-gather hooks are schedule-invariant across both FSDP2 passes, and + regression-guards the future P1.1 buffer-split bandwidth optimization. + """ + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size, device = _get_dist_info() + + # Shared, fixed input/target so the two schedules are compared on identical data. + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + def assert_exact(actual, expected, path): + if torch.is_tensor(expected): + torch.testing.assert_close( + actual, + expected, + rtol=0.0, + atol=0.0, + msg=lambda m: f"reshard schedule changed {path}: {m}", + ) + elif isinstance(expected, dict): + assert actual.keys() == expected.keys(), f"{path}: keys differ" + for key in expected: + assert_exact(actual[key], expected[key], f"{path}.{key}") + elif isinstance(expected, (list, tuple)): + assert len(actual) == len(expected), f"{path}: lengths differ" + for index, (actual_item, expected_item) in enumerate(zip(actual, expected)): + assert_exact(actual_item, expected_item, f"{path}[{index}]") + else: + assert actual == expected, f"{path}: {actual!r} != {expected!r}" + + def run(reshard_after_forward): + # Re-seed so both schedules get identical weight init from reset_parameters(). + torch.manual_seed(42) + torch.cuda.manual_seed(42) + model = _shard_model( + _build_hybrid_model(hybrid_recipe), + world_size, + reshard_after_forward=reshard_after_forward, + ) + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + run_x = x.detach().clone().requires_grad_() + artifacts = { + "outputs": [], + "losses": [], + "input_grads": [], + "param_grads": [], + "optimizer": [], + } + for _ in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + run_x.grad = None + with te.autocast(enabled=True, recipe=hybrid_recipe): + output = model(run_x) + artifacts["outputs"].append(output.detach().clone()) + loss = F.mse_loss(output, target) + artifacts["losses"].append(loss.detach().clone()) + loss.backward() + artifacts["input_grads"].append(run_x.grad.detach().clone()) + step_grads = [] + for param in model.parameters(): + grad = param.grad + if grad is not None: + grad = grad.to_local() if isinstance(grad, DTensor) else grad + grad = grad.detach().clone() + step_grads.append(grad) + artifacts["param_grads"].append(step_grads) + optimizer.step() + step_state = [] + for param in model.parameters(): + param_state = {} + for key, value in optimizer.state[param].items(): + if torch.is_tensor(value): + value = value.to_local() if isinstance(value, DTensor) else value + value = value.detach().clone() + param_state[key] = value + step_state.append(param_state) + artifacts["optimizer"].append(step_state) + return artifacts + + artifacts_resharded = run(reshard_after_forward=True) # re-gather in backward + artifacts_kept = run(reshard_after_forward=False) # keep gathered weight through backward + + losses_resharded = [loss.item() for loss in artifacts_resharded["losses"]] + assert all( + losses_resharded[i + 1] < losses_resharded[i] for i in range(NUM_STEPS - 1) + ), f"reshard_after_forward=True loss not strictly decreasing: {losses_resharded}" + assert_exact(artifacts_resharded, artifacts_kept, "training") + + +def test_fused_adam_hybrid_bf16_vs_hybrid_parity(hybrid_recipe_name): + """Compare hybrid+FSDP2 loss trajectory against BF16+FSDP2 within tolerance. + + This is a sanity check that hybrid quantized training converges similarly + to BF16 training, not a bitwise-exact comparison. + """ + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size, device = _get_dist_info() + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + def run_training(model, recipe_for_autocast): + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + losses = [] + for _ in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=(recipe_for_autocast is not None), recipe=recipe_for_autocast): + output = model(x) + loss = F.mse_loss(output, target) + losses.append(loss.item()) + loss.backward() + optimizer.step() + return losses + + # BF16 baseline + torch.manual_seed(42) + torch.cuda.manual_seed(42) + bf16_model = _build_model(fp8_init=False) + bf16_model = _shard_model(bf16_model, world_size) + bf16_losses = run_training(bf16_model, None) + + # Hybrid + torch.manual_seed(42) + torch.cuda.manual_seed(42) + hybrid_model = _build_hybrid_model(hybrid_recipe) + hybrid_model = _shard_model(hybrid_model, world_size) + hybrid_losses = run_training(hybrid_model, hybrid_recipe) + + assert hybrid_losses[-1] < hybrid_losses[0], f"Hybrid loss did not decrease: {hybrid_losses}" + assert bf16_losses[-1] < bf16_losses[0], f"BF16 loss did not decrease: {bf16_losses}" + + # Hybrid stays within a few % of bf16 (seed-fixed). + rel_tol = 0.10 + for step, (h_loss, b_loss) in enumerate(zip(hybrid_losses, bf16_losses)): + rel_diff = abs(h_loss - b_loss) / max(abs(b_loss), 1e-10) + assert rel_diff < rel_tol, ( + f"Step {step}: hybrid loss ({h_loss:.4f}) vs bf16 ({b_loss:.4f}) " + f"differ by {rel_diff * 100:.2f}% (> {rel_tol * 100:.0f}%)" + ) + + +# Same-format hybrid -> the vanilla recipe it must match bitwise. Cross-format +# hybrids (e.g. HybridMixed_MXFP8_FP8) have no single-format vanilla equivalent. +_HYBRID_TO_BASE_RECIPE = { + "HybridFP8CurrentScaling": "Float8CurrentScaling", + "HybridMXFP8": "MXFP8BlockScaling", + "HybridFloat8BlockScaling": "Float8BlockScaling", +} + + +def _build_linear_parity_stack(recipe): + """Two bare ``te.Linear`` layers under ``quantized_model_init`` for + hybrid-vs-vanilla bitwise parity. + """ + with te.quantized_model_init(enabled=True, recipe=recipe): + return torch.nn.Sequential( + te.Linear(HIDDEN_SIZE, HIDDEN_SIZE, params_dtype=torch.bfloat16, device="meta"), + te.Linear(HIDDEN_SIZE, HIDDEN_SIZE, params_dtype=torch.bfloat16, device="meta"), + ) + + +def test_fused_adam_hybrid_vs_base_recipe_parity(hybrid_recipe_name): + """Same-format hybrid must match its vanilla recipe bitwise through the full + FSDP2 + FusedAdam loop. + + Every output, loss, input/parameter gradient, master parameter, optimizer + moment, and step counter is asserted bitwise-identical -- a regression guard + for both amax reduction and master-weight requantization. Uses a bare + ``te.Linear`` stack (see ``_build_linear_parity_stack``) to isolate + GEMM-operand quantization. + """ + if hybrid_recipe_name not in _HYBRID_TO_BASE_RECIPE: + pytest.skip( + f"{hybrid_recipe_name} is cross-format; no single-format vanilla " + "recipe to compare against." + ) + + from fsdp2_utils import get_hybrid_recipe_from_string + + base_recipe_name = _HYBRID_TO_BASE_RECIPE[hybrid_recipe_name] + world_size, device = _get_dist_info() + + # Shared, fixed input/target; the comparison is per-rank (base vs hybrid on + # the same rank), so cross-rank input consistency does not matter. + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + def run_training(build_fn, recipe_for_autocast): + def snapshot_optimizer(model, optimizer): + snapshots = [] + for param in model.parameters(): + state = {} + for key, value in optimizer.state[param].items(): + if torch.is_tensor(value): + value = value.to_local() if isinstance(value, DTensor) else value + value = value.detach().clone() + state[key] = value + snapshots.append(state) + return snapshots + + # Re-seed so both models get identical init from reset_parameters() (run + # after sharding); with same-format quantization and a dropout-free loop + # the full trajectory is then deterministic. + torch.manual_seed(1234) + torch.cuda.manual_seed(1234) + model = _shard_model(build_fn(), world_size) + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + run_x = x.detach().clone().requires_grad_() + outputs = [] + losses = [] + input_grads = [] + grads_per_step = [] + optimizer_states = [] + for step in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + run_x.grad = None + with te.autocast(enabled=True, recipe=recipe_for_autocast): + output = model(run_x) + outputs.append(output.detach().clone()) + loss = F.mse_loss(output, target) + losses.append(loss.detach().clone()) + loss.backward() + input_grads.append(run_x.grad.detach().clone()) + # Snapshot grad local shards before the optimizer consumes them + # (p.grad is a DTensor under FSDP2) to assert backward parity directly. + step_grads = [] + for p in model.parameters(): + g = p.grad + if g is None: + step_grads.append(None) + else: + g = g.to_local() if isinstance(g, DTensor) else g + step_grads.append(g.detach().clone()) + grads_per_step.append(step_grads) + optimizer.step() + optimizer_states.append(snapshot_optimizer(model, optimizer)) + return outputs, losses, input_grads, grads_per_step, optimizer_states + + base_recipe = get_recipe_from_string(base_recipe_name) + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + + base_outputs, base_losses, base_input_grads, base_grads, base_opt_states = run_training( + lambda: _build_linear_parity_stack(base_recipe), base_recipe + ) + ( + hybrid_outputs, + hybrid_losses, + hybrid_input_grads, + hybrid_grads, + hybrid_opt_states, + ) = run_training(lambda: _build_linear_parity_stack(hybrid_recipe), hybrid_recipe) + + # (1) Every forward: bitwise-identical before and after optimizer updates. + for step, (base_output, hybrid_output) in enumerate(zip(base_outputs, hybrid_outputs)): + torch.testing.assert_close( + hybrid_output, + base_output, + rtol=0.0, + atol=0.0, + msg=lambda m, s=step: f"[{hybrid_recipe_name} vs {base_recipe_name}] step {s} forward output not bitwise-identical: {m}", + ) + + # (2) Every per-step loss: bitwise-identical across the whole optimizer loop. + for step, (b_loss, h_loss) in enumerate(zip(base_losses, hybrid_losses)): + torch.testing.assert_close( + h_loss, + b_loss, + rtol=0.0, + atol=0.0, + msg=lambda m, s=step: f"[{hybrid_recipe_name} vs {base_recipe_name}] step {s} loss not bitwise-identical to the vanilla recipe: {m}", + ) + + # (3) Backward: every weight-gradient shard at every step bitwise-identical + for step, (base_grad, hybrid_grad) in enumerate(zip(base_input_grads, hybrid_input_grads)): + torch.testing.assert_close( + hybrid_grad, + base_grad, + rtol=0.0, + atol=0.0, + msg=lambda m, s=step: f"[{hybrid_recipe_name} vs {base_recipe_name}] step {s} input gradient not bitwise-identical: {m}", + ) + + # (implied by the loss trajectory, but asserted directly to be explicit). + for step, (b_step, h_step) in enumerate(zip(base_grads, hybrid_grads)): + for i, (b_grad, h_grad) in enumerate(zip(b_step, h_step)): + assert (b_grad is None) == (h_grad is None), ( + f"[{hybrid_recipe_name} vs {base_recipe_name}] step {step} param {i}" + " gradient presence differs between hybrid and vanilla" + ) + if b_grad is None: + continue + torch.testing.assert_close( + h_grad, + b_grad, + rtol=0.0, + atol=0.0, + msg=lambda m, s=step, i=i: f"[{hybrid_recipe_name} vs {base_recipe_name}] step {s} param {i} gradient not bitwise-identical to the vanilla recipe: {m}", + ) + + # Optimizer continuation state: FP32 master params, moments, and counters. + assert len(base_opt_states) == len(hybrid_opt_states) + for step, (base_step, hybrid_step) in enumerate(zip(base_opt_states, hybrid_opt_states)): + assert len(base_step) == len(hybrid_step) + for param_idx, (base_state, hybrid_state) in enumerate(zip(base_step, hybrid_step)): + assert base_state.keys() == hybrid_state.keys(), ( + f"[{hybrid_recipe_name} vs {base_recipe_name}] step {step} param {param_idx} " + f"optimizer state keys differ: {base_state.keys()} != {hybrid_state.keys()}" + ) + for key in base_state: + base_value = base_state[key] + hybrid_value = hybrid_state[key] + if torch.is_tensor(base_value): + torch.testing.assert_close( + hybrid_value, + base_value, + rtol=0.0, + atol=0.0, + msg=lambda m, s=step, i=param_idx, k=key: f"[{hybrid_recipe_name} vs {base_recipe_name}] step {s} param {i} optimizer state {k!r} not bitwise-identical: {m}", + ) + else: + assert hybrid_value == base_value, ( + f"[{hybrid_recipe_name} vs {base_recipe_name}] step {step} " + f"param {param_idx} optimizer state {key!r} differs: " + f"{hybrid_value!r} != {base_value!r}" + ) + + +def test_fused_adam_hybrid_scale_uniform_across_shards(hybrid_recipe_name): + """Per-tensor hybrid weights must share ONE amax-reduced scale across FSDP2 + shards -- tolerance-free regression guard for the amax-reduction fix. + + Without cross-shard reduction each rank quantizes its shard with a local amax + and the scales differ; with the fix they match. Checked directly on the + sharded weight (no forward). Block-scaled formats (MXFP8) are skipped. + """ + if hybrid_recipe_name != "HybridFP8CurrentScaling": + pytest.skip("scale-uniformity check applies to per-tensor current scaling only") + + from transformer_engine.pytorch import HybridQuantizedTensor + from fsdp2_utils import get_hybrid_recipe_from_string + + world_size, device = _get_dist_info() + if world_size < 2: + pytest.skip("needs >=2 ranks to compare shard scales") + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + model = _build_hybrid_model(hybrid_recipe) + model = _shard_model(model, world_size) + + checked = {"rowwise": 0, "columnwise": 0} + for name, param in model.named_parameters(): + if not ( + isinstance(param, DTensor) and isinstance(param._local_tensor, HybridQuantizedTensor) + ): + continue + for direction in checked: + sub_storage = getattr(param._local_tensor, f"_{direction}_storage") + scale_inv = getattr(sub_storage, "_scale_inv", None) + if scale_inv is None: + continue + local_scale = scale_inv.detach().reshape(-1).clone() + gathered = [torch.zeros_like(local_scale) for _ in range(world_size)] + dist.all_gather(gathered, local_scale) + for r in range(1, world_size): + torch.testing.assert_close( + gathered[r], + gathered[0], + rtol=0.0, + atol=0.0, + msg=lambda m, n=name, r=r, d=direction: f"{n}: rank {r} {d} _scale_inv differs from rank 0 -- cross-shard amax reduction was not applied to the hybrid current-scaling weight: {m}", + ) + checked[direction] += 1 + assert all( + count > 0 for count in checked.values() + ), f"missing hybrid current-scaling directions: {checked}" + + +def test_fused_adam_hybrid_identity_fp8_master_weights(): + """FSDP2 + FusedAdam with Hybrid(FP8 current rowwise, Identity columnwise). + + Covers the Identity sub-storage in hybrid FSDP2 all-gather while the FP8 + current rowwise direction validates cross-shard amax reduction via scale + uniformity. + """ + from transformer_engine.pytorch import HybridQuantizedTensor + from transformer_engine.pytorch.tensor.storage.identity_tensor_storage import ( + IdentityTensorStorage, + ) + from fsdp2_utils import get_hybrid_recipe_from_string + + world_size, device = _get_dist_info() + if world_size < 2: + pytest.skip("needs >=2 ranks to validate cross-shard amax reduction") + + hybrid_recipe = get_hybrid_recipe_from_string("HybridFP8CurrentScalingIdentity") + model = _build_linear_parity_stack(hybrid_recipe) + model = _shard_model(model, world_size) + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + losses = [] + identity_steps = 2 + for step in range(identity_steps): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + loss = F.mse_loss(output, target) + losses.append(loss.item()) + loss.backward() + if step < identity_steps - 1: + optimizer.step() + + assert all( + losses[i + 1] < losses[i] for i in range(len(losses) - 1) + ), f"Hybrid Identity/FP8 loss not strictly decreasing: {losses}" + + checked_identity = 0 + checked_scale = 0 + for name, param in model.named_parameters(): + if not ( + isinstance(param, DTensor) and isinstance(param._local_tensor, HybridQuantizedTensor) + ): + continue + local = param._local_tensor + assert isinstance(local._columnwise_storage, IdentityTensorStorage) + + scale_inv = getattr(local._rowwise_storage, "_scale_inv", None) + if scale_inv is not None: + local_scale = scale_inv.detach().reshape(-1).clone() + gathered_scales = [torch.zeros_like(local_scale) for _ in range(world_size)] + dist.all_gather(gathered_scales, local_scale) + for r in range(1, world_size): + torch.testing.assert_close( + gathered_scales[r], + gathered_scales[0], + rtol=0.0, + atol=0.0, + msg=lambda m, n=name, r=r: f"{n}: rank {r} rowwise _scale_inv differs from rank 0 for Hybrid(FP8Current, Identity): {m}", + ) + checked_scale += 1 + + local_identity = local._columnwise_storage.dequantize().contiguous() + gathered_identity = [torch.zeros_like(local_identity) for _ in range(world_size)] + dist.all_gather(gathered_identity, local_identity) + manual_full = torch.cat(gathered_identity, dim=0) + + sharded_tensors, metadata = local.fsdp_pre_all_gather( + mesh=None, + orig_size=local.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + all_gather_outputs = [] + for shard in sharded_tensors: + gathered = [torch.zeros_like(shard) for _ in range(world_size)] + dist.all_gather(gathered, shard) + all_gather_outputs.append(torch.cat(gathered, dim=0)) + fsdp_full, _ = local.fsdp_post_all_gather( + tuple(all_gather_outputs), metadata, local.dtype, out=None + ) + assert isinstance(fsdp_full, HybridQuantizedTensor) + full_identity = fsdp_full._columnwise_storage.dequantize() + torch.testing.assert_close( + manual_full.float(), + full_identity[: manual_full.shape[0]].float(), + rtol=0.0, + atol=0.0, + msg=lambda m, n=name: f"{n}: Identity columnwise all-gather mismatch: {m}", + ) + checked_identity += 1 + + assert checked_identity > 0, "no Hybrid(FP8Current, Identity) params found" + assert checked_scale > 0, "no FP8 current rowwise scales found to check" + + +def test_fused_adam_hybrid_allgather_correctness(hybrid_recipe_name): + """Validate both Hybrid directions through FSDP2 at raw-buffer precision. + + The rowwise and columnwise sub-storages are checked independently. Every + gathered data buffer, scale buffer, field layout, and per-tensor scale + metadata value must exactly match a manual dim-0 all-gather. + """ + from fsdp2_utils import get_hybrid_recipe_from_string + + expected_types = { + "HybridFP8CurrentScaling": (Float8TensorStorage, Float8TensorStorage), + "HybridMXFP8": (MXFP8TensorStorage, MXFP8TensorStorage), + "HybridFloat8BlockScaling": ( + Float8BlockwiseQTensorStorage, + Float8BlockwiseQTensorStorage, + ), + "HybridMixed_MXFP8_FP8": (MXFP8TensorStorage, Float8TensorStorage), + } + row_type, col_type = expected_types[hybrid_recipe_name] + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size, device = _get_dist_info() + + model = _build_hybrid_model(hybrid_recipe) + model = _shard_model(model, world_size) + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + with te.autocast(enabled=True, recipe=hybrid_recipe): + _ = model(x) + + params = [ + (name, param) + for name, param in model.named_parameters() + if isinstance(param, DTensor) and isinstance(param._local_tensor, HybridQuantizedTensor) + ] + local_count = torch.tensor([len(params)], dtype=torch.int64, device=device) + min_count = local_count.clone() + max_count = local_count.clone() + dist.all_reduce(min_count, dist.ReduceOp.MIN) + dist.all_reduce(max_count, dist.ReduceOp.MAX) + _collective_assert( + min_count.item() == max_count.item() and min_count.item() > 0, + f"{hybrid_recipe_name}: Hybrid parameter count range is " + f"({min_count.item()}, {max_count.item()})", + ) + + errors = [] + for name, param in params: + local = param._local_tensor + reconstructed = _manual_reconstruct_hybrid(local, param_name=name, world_size=world_size) + public_full = param.full_tensor() + _record_exact( + errors, + public_full, + reconstructed.dequantize(), + f"{name}: public full_tensor vs Hybrid reconstruction", + ) + _check_hybrid_direction_buffers( + local._rowwise_storage, + reconstructed._rowwise_storage, + row_type, + direction="rowwise", + param_name=name, + world_size=world_size, + errors=errors, + ) + _check_hybrid_direction_buffers( + local._columnwise_storage, + reconstructed._columnwise_storage, + col_type, + direction="columnwise", + param_name=name, + world_size=world_size, + errors=errors, + ) + + _raise_collective_errors(errors, f"{hybrid_recipe_name} raw Hybrid all-gather") + + +def test_fused_adam_hybrid_mxfp8_awkward_shard_shape(): + """Exercise MXFP8 block-scale unpad/pad on a sharded Linear whose shard + dim-0 is block-aligned (divisible by 32) but NOT divisible by 128. + + MXFP8 block scales are stored with ``[128, 4]`` / ``[4, 128]`` alignment + padding, which must be stripped before FSDP2's dim-0 all-gather and + re-applied after. With ``HIDDEN_SIZE`` and ``FFN_HIDDEN_SIZE`` both + divisible by 128, the default model never forces this code path, so this + test uses a hand-picked Linear size. + + Regression test for the "pre-fix" bug where + ``HybridQuantizedTensor.fsdp_pre_all_gather`` pulled raw tensor fields via + ``get_metadata()`` without unpadding the scale — the padded bytes would + have been interleaved at every rank boundary in the gather output. + """ + from fsdp2_utils import get_hybrid_recipe_from_string + + supported, reason = te.is_mxfp8_available(return_reason=True) + if not supported: + pytest.skip(f"MXFP8: {reason}") + + world_size, device = _get_dist_info() + + # FSDP2 shards a Linear weight of shape (out_features, in_features) along + # dim-0, so each rank holds `out_features / world_size` rows. Pick + # per-rank shard dim-0 = 96: divisible by MXFP8_BLOCK_SCALING_SIZE (32) + # so data alignment holds, but NOT divisible by 128 so the rowwise + # scale-inv needs alignment padding on the sharded copy. This is the + # shape that exercises the unpad-before-gather / pad-after-gather + # behaviour in MXFP8TensorStorage.fsdp_{extract,assign}_buffers. + per_rank_out = 96 + out_features = per_rank_out * world_size + in_features = 128 # arbitrary, divisible by 32; not sharded by FSDP2 here + assert per_rank_out % 32 == 0, ( + f"Test setup error: per_rank_out={per_rank_out} (= out_features / world_size, " + f"world_size={world_size}) must be a multiple of the MXFP8 block size (32) so the " + "sharded weight's data stays block-aligned. Pick a per_rank_out divisible by 32." + ) + assert per_rank_out % 128 != 0, ( + f"Test setup error: per_rank_out={per_rank_out} must NOT be a multiple of 128, or the " + "rowwise scale-inv needs no alignment padding and this test stops exercising the MXFP8 " + "unpad-before-gather / pad-after-gather path it exists to cover. Pick a per_rank_out " + "divisible by 32 but not 128 (e.g. 96)." + ) + + for recipe_name in ("HybridMXFP8", "HybridMixed_MXFP8_FP8"): + hybrid_recipe = get_hybrid_recipe_from_string(recipe_name) + + with te.quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = torch.nn.Sequential( + te.Linear( + in_features, + out_features, + params_dtype=torch.bfloat16, + device="meta", + ), + ) + model = _shard_model(model, world_size) + + # Batch (leading) dim must be divisible by MXFP8_BLOCK_SCALING_SIZE (32). + x = torch.randn(32, in_features, dtype=torch.bfloat16, device=device) + with te.autocast(enabled=True, recipe=hybrid_recipe): + out = model(x) + out.sum().backward() + + col_type = MXFP8TensorStorage if recipe_name == "HybridMXFP8" else Float8TensorStorage + params = [ + (name, param) + for name, param in model.named_parameters() + if isinstance(param, DTensor) and isinstance(param._local_tensor, HybridQuantizedTensor) + ] + local_count = torch.tensor([len(params)], dtype=torch.int64, device=device) + min_count = local_count.clone() + max_count = local_count.clone() + dist.all_reduce(min_count, dist.ReduceOp.MIN) + dist.all_reduce(max_count, dist.ReduceOp.MAX) + _collective_assert( + min_count.item() == max_count.item() and min_count.item() > 0, + f"{recipe_name}: awkward-shape Hybrid parameter count range is " + f"({min_count.item()}, {max_count.item()})", + ) + + errors = [] + for name, param in params: + local = param._local_tensor + label = f"{recipe_name}:{name}" + reconstructed = _manual_reconstruct_hybrid( + local, param_name=label, world_size=world_size + ) + public_full = param.full_tensor() + _record_exact( + errors, + public_full, + reconstructed.dequantize(), + f"{label}: public full_tensor vs Hybrid reconstruction", + ) + _check_hybrid_direction_buffers( + local._rowwise_storage, + reconstructed._rowwise_storage, + MXFP8TensorStorage, + direction="rowwise", + param_name=label, + world_size=world_size, + errors=errors, + ) + _check_hybrid_direction_buffers( + local._columnwise_storage, + reconstructed._columnwise_storage, + col_type, + direction="columnwise", + param_name=label, + world_size=world_size, + errors=errors, + ) + _raise_collective_errors(errors, f"{recipe_name} awkward raw Hybrid all-gather") + + +def test_fused_adam_hybrid_float8_block_unaligned_shard_shape(): + """Unaligned local shards are rejected before gathering incompatible scale tiles.""" + from transformer_engine.pytorch import fp8 + from fsdp2_utils import get_hybrid_recipe_from_string + + supported, reason = fp8.check_fp8_block_scaling_support() + if not supported: + pytest.skip(reason) + + world_size, device = _get_dist_info() + if world_size != 2: + pytest.skip("test shape is defined for exactly two FSDP ranks") + + in_features = 256 + per_rank_out = 192 + out_features = per_rank_out * world_size + assert out_features % 128 == 0 + assert per_rank_out % 128 != 0 + + hybrid_recipe = get_hybrid_recipe_from_string("HybridFloat8BlockScaling") + with te.quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = torch.nn.Sequential( + te.Linear( + in_features, + out_features, + params_dtype=torch.bfloat16, + device="meta", + ) + ) + model = _shard_model(model, world_size) + + x = torch.randn(128, in_features, dtype=torch.bfloat16, device=device) + with pytest.raises( + RuntimeError, + match="local flattened M dimension.*not a multiple of 128", + ): + with te.autocast(enabled=True, recipe=hybrid_recipe): + model(x) + + +def test_hybrid_dcp_output_parity(hybrid_recipe_name): + """DCP roundtrip and exact forked optimizer continuation. + + Trains and checkpoints a hybrid model plus FusedAdam, loads both into fresh + objects, and compares the identical next step against uninterrupted training: + output, loss, input/parameter gradients, master params, moments, counters, + and post-step output must all match bitwise. + """ + import torch.distributed.checkpoint as dcp + + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size, device = _get_dist_info() + rank = int(os.environ.get("RANK", "0")) + # Deterministic, rank-agnostic checkpoint dir so all ranks read/write + # the same DCP path. ``os.getpid()`` differs per rank under torchrun. + checkpoint_dir = f"/tmp/te_test_fsdp2_hybrid_dcp_parity_{hybrid_recipe_name}" + + if rank == 0: + shutil.rmtree(checkpoint_dir, ignore_errors=True) + dist.barrier() + + try: + model = _build_hybrid_model(hybrid_recipe) + model = _shard_model(model, world_size) + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + failures = [] + + def snapshot_optimizer(current_model, current_optimizer): + snapshots = [] + for param in current_model.parameters(): + state = {} + for key, value in current_optimizer.state[param].items(): + if torch.is_tensor(value): + value = value.to_local() if isinstance(value, DTensor) else value + value = value.detach().clone() + state[key] = value + snapshots.append(state) + return snapshots + + def check_optimizer_state(actual, expected, label): + if len(actual) != len(expected): + failures.append( + f"{label}: parameter count differs: {len(actual)} != {len(expected)}" + ) + for param_idx, (actual_state, expected_state) in enumerate(zip(actual, expected)): + if actual_state.keys() != expected_state.keys(): + failures.append( + f"{label}: param {param_idx} state keys differ: " + f"{actual_state.keys()} != {expected_state.keys()}" + ) + for key in actual_state.keys() & expected_state.keys(): + actual_value = actual_state[key] + expected_value = expected_state[key] + if torch.is_tensor(expected_value): + _record_exact( + failures, + actual_value, + expected_value, + f"{label}: param {param_idx} optimizer state {key!r}", + ) + elif actual_value != expected_value: + failures.append( + f"{label}: param {param_idx} optimizer state {key!r} differs: " + f"{actual_value!r} != {expected_value!r}" + ) + + def run_continuation_step(current_model, current_optimizer): + step_x = x.detach().clone().requires_grad_() + current_optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=hybrid_recipe): + step_output = current_model(step_x) + step_loss = F.mse_loss(step_output, target) + step_loss.backward() + step_grads = [] + for param in current_model.parameters(): + grad = param.grad + if grad is not None: + grad = grad.to_local() if isinstance(grad, DTensor) else grad + grad = grad.detach().clone() + step_grads.append(grad) + current_optimizer.step() + with torch.no_grad(), te.autocast(enabled=True, recipe=hybrid_recipe): + post_step_output = current_model(x).detach().clone() + return { + "output": step_output.detach().clone(), + "loss": step_loss.detach().clone(), + "input_grad": None if step_x.grad is None else step_x.grad.detach().clone(), + "param_grads": step_grads, + "optimizer": snapshot_optimizer(current_model, current_optimizer), + "post_step_output": post_step_output, + } + + for _ in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + F.mse_loss(output, target).backward() + optimizer.step() + + with torch.no_grad(): + with te.autocast(enabled=True, recipe=hybrid_recipe): + ref_output = model(x).clone() + + save_state = { + "model": model.state_dict(), + "optimizer": optimizer.state_dict(), + } + dcp.save(save_state, checkpoint_id=checkpoint_dir) + saved_optimizer_state = snapshot_optimizer(model, optimizer) + + model2 = _build_hybrid_model(hybrid_recipe) + model2 = _shard_model(model2, world_size) + optimizer2 = te.optimizers.FusedAdam( + model2.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + optimizer2.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=hybrid_recipe): + out_tmp = model2(x) + F.mse_loss(out_tmp, target).backward() + optimizer2.step() + + state_to_load = { + "model": model2.state_dict(), + "optimizer": optimizer2.state_dict(), + } + dcp.load(state_to_load, checkpoint_id=checkpoint_dir) + model2.load_state_dict(state_to_load["model"]) + optimizer2.load_state_dict(state_to_load["optimizer"]) + + with torch.no_grad(): + with te.autocast(enabled=True, recipe=hybrid_recipe): + loaded_output = model2(x) + check_optimizer_state( + snapshot_optimizer(model2, optimizer2), + saved_optimizer_state, + "DCP optimizer state immediately after load", + ) + _record_exact(failures, loaded_output, ref_output, "DCP roundtrip output") + _raise_collective_errors(failures, "DCP state immediately after load") + failures.clear() + + # Fork from the checkpoint and execute the identical next step. This + # catches an ignored optimizer checkpoint even when model-only output + # parity succeeds immediately after load. + torch.manual_seed(9876) + torch.cuda.manual_seed(9876) + reference_step = run_continuation_step(model, optimizer) + torch.manual_seed(9876) + torch.cuda.manual_seed(9876) + resumed_step = run_continuation_step(model2, optimizer2) + + for key in ("output", "loss", "input_grad", "post_step_output"): + _record_exact( + failures, + resumed_step[key], + reference_step[key], + f"DCP continuation {key}", + ) + + resumed_grads = resumed_step["param_grads"] + reference_grads = reference_step["param_grads"] + if len(resumed_grads) != len(reference_grads): + failures.append( + "DCP continuation parameter-gradient count differs: " + f"{len(resumed_grads)} != {len(reference_grads)}" + ) + for param_idx, (resumed_grad, reference_grad) in enumerate( + zip(resumed_grads, reference_grads) + ): + _record_exact( + failures, + resumed_grad, + reference_grad, + f"DCP continuation param {param_idx} gradient", + ) + + check_optimizer_state( + resumed_step["optimizer"], + reference_step["optimizer"], + "DCP optimizer state after continuation", + ) + _raise_collective_errors(failures, "DCP forked continuation") + finally: + dist.barrier() + if rank == 0: + shutil.rmtree(checkpoint_dir, ignore_errors=True) + + TESTS = { "fused_adam_fp8_master_weights": test_fused_adam_fp8_master_weights, "fused_adam_fp8_master_weights_no_meta": test_fused_adam_fp8_master_weights_no_meta, @@ -1140,6 +2376,24 @@ def test_dcp_resharding_load(recipe_name): "dcp_resharding_save": test_dcp_resharding_save, "dcp_resharding_load": test_dcp_resharding_load, "safetensors_fp32_export": test_safetensors_fp32_export, + "fused_adam_hybrid_master_weights": test_fused_adam_hybrid_master_weights, + "fused_adam_hybrid_bf16_vs_hybrid_parity": test_fused_adam_hybrid_bf16_vs_hybrid_parity, + "fused_adam_hybrid_vs_base_recipe_parity": test_fused_adam_hybrid_vs_base_recipe_parity, + "fused_adam_hybrid_scale_uniform_across_shards": ( + test_fused_adam_hybrid_scale_uniform_across_shards + ), + "fused_adam_hybrid_identity_fp8_master_weights": ( + test_fused_adam_hybrid_identity_fp8_master_weights + ), + "fused_adam_hybrid_allgather_correctness": test_fused_adam_hybrid_allgather_correctness, + "fused_adam_hybrid_mxfp8_awkward_shard_shape": test_fused_adam_hybrid_mxfp8_awkward_shard_shape, + "hybrid_dcp_output_parity": test_hybrid_dcp_output_parity, +} + +# Hybrid tests that are NOT parametrized by recipe (they sweep internally). +_HYBRID_NON_PARAMETRIZED_TESTS = { + "fused_adam_hybrid_identity_fp8_master_weights", + "fused_adam_hybrid_mxfp8_awkward_shard_shape", } @@ -1157,6 +2411,11 @@ def test_dcp_resharding_load(recipe_name): "Float8BlockScaling", "MXFP8BlockScaling", "NVFP4BlockScaling", + "HybridFP8CurrentScaling", + "HybridMXFP8", + "HybridFloat8BlockScaling", + "HybridMixed_MXFP8_FP8", + "HybridFP8CurrentScalingIdentity", ], ) args = parser.parse_args() @@ -1166,7 +2425,10 @@ def test_dcp_resharding_load(recipe_name): torch.manual_seed(42) torch.cuda.manual_seed(42) try: - TESTS[args.test](args.recipe) + if args.test in _HYBRID_NON_PARAMETRIZED_TESTS: + TESTS[args.test]() + else: + TESTS[args.test](args.recipe) finally: # NOTE: In PyTorch < 2.6 there’s a teardown race where one rank may call # destroy_process_group() while other ranks still have in-flight NCCL ops, diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py index 81c4d3e888..8a015fe659 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py @@ -469,12 +469,169 @@ def test_transpose_cache_retained_after_backward(recipe_name, quantized_model_in ) +# ── Hybrid quantization memory tests ───────────────────────────────── + + +def _build_hybrid_model(num_layers, hybrid_recipe, use_meta_device=True): + """Build a model with quantized_model_init using a hybrid CustomRecipe.""" + kwargs = dict( + fuse_qkv_params=True, + params_dtype=torch.bfloat16, + hidden_dropout=0.0, + attention_dropout=0.0, + ) + if use_meta_device: + kwargs["device"] = "meta" + with te.quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = torch.nn.Sequential( + *[ + te.TransformerLayer( + HIDDEN_SIZE, + FFN_HIDDEN_SIZE, + NUM_ATTENTION_HEADS, + **kwargs, + ) + for _ in range(num_layers) + ] + ) + return model + + +def test_hybrid_no_excess_forward_memory(hybrid_recipe_name): + """Hybrid quantized weights should not accumulate across layers during forward. + + Same methodology as test_fp8_temp_accumulation_across_layers but for + hybrid quantized tensors. + """ + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size, device = _get_dist_info() + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + # bf16 baseline + bf16_model = _build_model(NUM_LAYERS, fp8_init=False) + bf16_model = _shard_model(bf16_model, world_size) + bf16_optimizer = te.optimizers.FusedAdam( + bf16_model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + for _ in range(WARMUP_STEPS): + _run_training_step(bf16_model, bf16_optimizer, None, x, target) + bf16_increments = _measure_forward_increments(bf16_model, bf16_optimizer, None, x, target) + bf16_avg = sum(bf16_increments) / len(bf16_increments) + + del bf16_model, bf16_optimizer + gc.collect() + torch.cuda.empty_cache() + + # Hybrid model + hybrid_model = _build_hybrid_model(NUM_LAYERS, hybrid_recipe) + hybrid_model = _shard_model(hybrid_model, world_size) + hybrid_optimizer = te.optimizers.FusedAdam( + hybrid_model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + for _ in range(WARMUP_STEPS): + _run_training_step(hybrid_model, hybrid_optimizer, hybrid_recipe, x, target) + hybrid_increments = _measure_forward_increments( + hybrid_model, + hybrid_optimizer, + hybrid_recipe, + x, + target, + ) + hybrid_avg = sum(hybrid_increments) / len(hybrid_increments) + + excess_per_layer = hybrid_avg - bf16_avg + # Basis: forward growth is constant per layer (no accumulation) for both bf16 and + # hybrid; the excess is just hybrid's extra per-layer quantized buffers. Measured + # excess: ~3 KiB (FP8 current) / ~7 KiB (mixed MXFP8+FP8) / ~12 KiB (MXFP8). A + # leaked layer's quantized weights would be hundreds of KiB, so 50 KiB sits above + # the real per-layer overhead and well below a leak. + tolerance_per_layer = 50 * 1024 # 50 KiB + + assert excess_per_layer <= tolerance_per_layer, ( + "Hybrid per-layer forward memory increment exceeds bf16 baseline by " + f"{excess_per_layer/1024:.1f} KiB/layer (tolerance: {tolerance_per_layer/1024:.1f} KiB). " + f"bf16 avg: {bf16_avg/1024:.1f} KiB/layer, hybrid avg: {hybrid_avg/1024:.1f} KiB/layer." + ) + + +def test_hybrid_transpose_cache_after_backward(hybrid_recipe_name): + """Detect transpose caches from hybrid sub-storages persisting after backward.""" + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size, device = _get_dist_info() + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + # bf16 baseline + bf16_model = _build_model(NUM_LAYERS, fp8_init=False) + bf16_model = _shard_model(bf16_model, world_size) + bf16_optimizer = te.optimizers.FusedAdam( + bf16_model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + for _ in range(WARMUP_STEPS): + _run_training_step(bf16_model, bf16_optimizer, None, x, target) + bf16_bwd_delta = _measure_backward_memory_delta(bf16_model, bf16_optimizer, None, x, target) + + del bf16_model, bf16_optimizer + gc.collect() + torch.cuda.empty_cache() + + # Hybrid model + hybrid_model = _build_hybrid_model(NUM_LAYERS, hybrid_recipe) + hybrid_model = _shard_model(hybrid_model, world_size) + hybrid_optimizer = te.optimizers.FusedAdam( + hybrid_model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + for _ in range(WARMUP_STEPS): + _run_training_step(hybrid_model, hybrid_optimizer, hybrid_recipe, x, target) + hybrid_bwd_delta = _measure_backward_memory_delta( + hybrid_model, + hybrid_optimizer, + hybrid_recipe, + x, + target, + ) + + excess = hybrid_bwd_delta - bf16_bwd_delta + # Basis: hybrid retains no more than bf16 after backward+step — measured excess is + # slightly negative (~-0.02..-0.09 MiB vs a ~2 MiB bf16 delta). The tolerance only + # absorbs allocator/measurement noise; a genuinely retained gathered weight or + # transpose cache would be MiB-scale (>> 256 KiB). + tolerance = 256 * 1024 # 256 KiB + + assert excess <= tolerance, ( + f"Hybrid backward retains {excess/1024**2:.2f} MiB more than bf16 baseline. " + f"bf16 backward delta: {bf16_bwd_delta/1024**2:.2f} MiB, " + f"hybrid backward delta: {hybrid_bwd_delta/1024**2:.2f} MiB." + ) + + # ── Standalone runner ──────────────────────────────────────────────── TESTS = { "bf16_no_excess_forward_memory": test_bf16_no_excess_forward_memory, "bf16_no_excess_backward_memory": test_bf16_no_excess_backward_memory, "fp8_temp_accumulation_across_layers": test_fp8_temp_accumulation_across_layers, "transpose_cache_retained_after_backward": test_transpose_cache_retained_after_backward, + "hybrid_no_excess_forward_memory": test_hybrid_no_excess_forward_memory, + "hybrid_transpose_cache_after_backward": test_hybrid_transpose_cache_after_backward, } if __name__ == "__main__": @@ -490,6 +647,10 @@ def test_transpose_cache_retained_after_backward(recipe_name, quantized_model_in "Float8BlockScaling", "MXFP8BlockScaling", "NVFP4BlockScaling", + "HybridFP8CurrentScaling", + "HybridMXFP8", + "HybridFloat8BlockScaling", + "HybridMixed_MXFP8_FP8", ], ) parser.add_argument("--quantized-model-init", action="store_true", default=False) @@ -505,11 +666,17 @@ def test_transpose_cache_retained_after_backward(recipe_name, quantized_model_in "fp8_temp_accumulation_across_layers", "transpose_cache_retained_after_backward", } + _HYBRID_PARAMETRIZED_TESTS = { + "hybrid_no_excess_forward_memory", + "hybrid_transpose_cache_after_backward", + } try: test_fn = TESTS[args.test] if args.test in _PARAMETRIZED_TESTS: test_fn(args.recipe, args.quantized_model_init) + elif args.test in _HYBRID_PARAMETRIZED_TESTS: + test_fn(args.recipe) else: test_fn() finally: diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py index 5eee186ef2..bd52e40d00 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py @@ -29,6 +29,7 @@ """ import gc +import math import os import sys import argparse @@ -220,10 +221,14 @@ def shard_model_with_fsdp2(model, mesh): @torch.no_grad() -def _check_fp8_fsdp2_allgather(model): - # Do manual allgather in fp32 and match against fp8 allgather done - # with fsdp2 - # FP32 manual weight allgather +def _check_fp8_fsdp2_allgather(model, *, tols=None): + """Compare FSDP2's quantized all-gather against a manual HP all-gather. + + ``tols=None`` preserves the format-specific defaults used by the original + FP8 checker. Callers for other quantized tensor compositions can provide a + single tolerance dictionary that applies to every parameter. + """ + # Manual high-precision weight all-gather. fp32_allgathered_params = {} for name, param in model.named_parameters(): assert isinstance(param, DTensor) @@ -234,29 +239,32 @@ def _check_fp8_fsdp2_allgather(model): if device_mesh.ndim > 1 else device_mesh.get_group() ) - # Perform manual allgather on local_tensor. zeros_like will create hp tensor since - # torch_dispatch for local_tensor will go down the dequantization route. + # Materialize high precision explicitly so this also works for composite + # quantized tensors such as HybridQuantizedTensor. + local_hp = local_tensor.dequantize() gathered_tensor = [ - torch.zeros_like(local_tensor) for _ in range(dist.get_world_size(group=dist_group)) + torch.zeros_like(local_hp) for _ in range(dist.get_world_size(group=dist_group)) ] - dist.all_gather(gathered_tensor, local_tensor.dequantize(), group=dist_group) + dist.all_gather(gathered_tensor, local_hp, group=dist_group) full_tensor = torch.cat(gathered_tensor, dim=0) fp32_allgathered_params[name] = full_tensor - # FP8 allgather using FSDP2 + # Quantized all-gather using FSDP2. for module in model.modules(): # Not all modules are wrapped/sharded with FSDP2. if hasattr(module, "unshard"): module.unshard() - # Make sure allgathered parameters match exactly + # Make sure all-gathered parameters match. for name, param in model.named_parameters(): - # NVFP4 scale unpad/repad through FSDP2 introduces small numerical - # differences vs the manual dequantize-then-allgather path. - if isinstance(param, NVFP4Tensor): - tols = dict(atol=5e-4, rtol=5e-3) + if tols is not None: + param_tols = tols + elif isinstance(param, NVFP4Tensor): + # NVFP4 scale unpad/repad through FSDP2 introduces small numerical + # differences vs the manual dequantize-then-allgather path. + param_tols = dict(atol=5e-4, rtol=5e-3) else: - tols = {} - torch.testing.assert_close(param.dequantize(), fp32_allgathered_params[name], **tols) - # Revert model to original sharded state + param_tols = {} + torch.testing.assert_close(param.dequantize(), fp32_allgathered_params[name], **param_tols) + # Revert model to original sharded state. for module in model.modules(): # Not all modules are wrapped/sharded with FSDP2. if hasattr(module, "reshard"): @@ -439,5 +447,254 @@ def test_distributed(recipe_name, fp8_init, sharding_dims, layer_type): _run_training(args) +def test_distributed_hybrid(hybrid_recipe_name): + """FSDP2 training with hybrid quantized_model_init. + + Uses quantized_model_init with a hybrid CustomRecipe on a TransformerLayer + model: + - params are DTensors wrapping HybridQuantizedTensor local shards, + - the loss is finite and strictly decreasing on fixed data (grads flow), + - params keep their hybrid quantized type across optimizer.step(), + - FSDP2's quantized all-gather matches a manual fp32 dequant-then-allgather. + """ + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + device = torch.device(f"cuda:{int(os.getenv('LOCAL_RANK', '0'))}") + + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + # Float8 block scaling requires every local FSDP shard's flattened M + # dimension to contain whole 128-row scale tiles. Keep the historical + # 512-wide model on up to four ranks, and scale it for larger worlds so + # the projection, fused-QKV, and MLP output dimensions all stay aligned. + shard_alignment = 128 * world_size + hidden_size = ((512 + shard_alignment - 1) // shard_alignment) * shard_alignment + ffn_hidden_size = 4 * hidden_size + assert hidden_size % shard_alignment == 0 + + kwargs = dict( + fuse_qkv_params=True, + params_dtype=torch.bfloat16, + hidden_dropout=0.0, + attention_dropout=0.0, + device="meta", + ) + with te.quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = torch.nn.Sequential( + *[te.TransformerLayer(hidden_size, ffn_hidden_size, 8, **kwargs) for _ in range(2)] + ) + + custom_attrs = save_custom_attrs(model) + mesh = get_device_mesh(world_size, [world_size]) + model = shard_model_with_fsdp2(model, mesh) + for module in model.modules(): + if hasattr(module, "reset_parameters"): + module.reset_parameters() + restore_custom_attrs(model, custom_attrs) + + from transformer_engine.pytorch import HybridQuantizedTensor + + def _hybrid_param_count(): + return sum( + 1 + for p in model.parameters() + if isinstance(p, DTensor) and isinstance(p._local_tensor, HybridQuantizedTensor) + ) + + # quantized_model_init must produce HybridQuantizedTensor local shards. + hybrid_count = _hybrid_param_count() + assert hybrid_count > 0, "No HybridQuantizedTensor local tensors after sharding" + + optimizer_lr = 1e-3 if hidden_size <= 512 else 1e-4 + optimizer = optim.Adam(model.parameters(), lr=optimizer_lr) + + input_data = torch.randn(128, 16, hidden_size, device=device, dtype=torch.bfloat16) + target = torch.randn(128, 16, hidden_size, device=device, dtype=torch.bfloat16) + + losses = [] + for iteration in range(3): + optimizer.zero_grad() + with te.autocast(enabled=True, recipe=hybrid_recipe): + output = model(input_data) + loss = F.mse_loss(output, target) + loss.backward() + optimizer.step() + loss_val = loss.item() + assert math.isfinite(loss_val), f"Non-finite loss at iter {iteration}: {loss_val}" + losses.append(loss_val) + dist_print(f"Hybrid iteration {iteration} completed with loss {loss_val}") + + # Training must actually progress on fixed data: strictly monotonic decrease. + assert all( + losses[i + 1] < losses[i] for i in range(len(losses) - 1) + ), f"Loss not strictly decreasing each step: {losses}" + + # Params must stay HybridQuantizedTensor after the optimizer step -- guards a + # silent dequantize-to-bf16 through the FSDP2 / optimizer path. + assert ( + _hybrid_param_count() == hybrid_count + ), "HybridQuantizedTensor params lost their quantized type after optimizer.step()" + + # FSDP2 quantized all-gather must match a manual fp32 dequant-then-allgather. + # Hybrid sub-storages (e.g. MXFP8 scale unpad/repad through FSDP2) can + # introduce small differences vs the manual dequantize-then-allgather path. + _check_fp8_fsdp2_allgather(model, tols=dict(atol=5e-4, rtol=5e-3)) + + +def test_distributed_hybrid_identity_all(): + """FSDP2 training/all-gather with an all-Identity CustomRecipe. + + This is the high-precision passthrough baseline for the hybrid/identity + tensor plumbing: quantized_model_init should produce IdentityTensor local + shards, optimizer steps should preserve that type, and FSDP2 all-gather + should reconstruct the same high-precision values as a manual gather. + """ + from transformer_engine.pytorch.tensor.identity_tensor import IdentityTensor + from fsdp2_utils import get_hybrid_recipe_from_string + + identity_recipe = get_hybrid_recipe_from_string("Identity") + world_size = int(os.environ.get("WORLD_SIZE", "1")) + device = torch.device(f"cuda:{int(os.getenv('LOCAL_RANK', '0'))}") + + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + kwargs = dict( + fuse_qkv_params=True, + params_dtype=torch.bfloat16, + hidden_dropout=0.0, + attention_dropout=0.0, + device="meta", + ) + with te.quantized_model_init(enabled=True, recipe=identity_recipe): + model = torch.nn.Sequential( + *[te.TransformerLayer(512, 2048, 8, **kwargs) for _ in range(2)] + ) + + custom_attrs = save_custom_attrs(model) + mesh = get_device_mesh(world_size, [world_size]) + model = shard_model_with_fsdp2(model, mesh) + for module in model.modules(): + if hasattr(module, "reset_parameters"): + module.reset_parameters() + restore_custom_attrs(model, custom_attrs) + + def _identity_param_count(): + return sum( + 1 + for p in model.parameters() + if isinstance(p, DTensor) and isinstance(p._local_tensor, IdentityTensor) + ) + + identity_count = _identity_param_count() + assert identity_count > 0, "No IdentityTensor local tensors after FSDP2 sharding" + + optimizer = optim.Adam(model.parameters(), lr=1e-3) + input_data = torch.randn(128, 16, 512, device=device, dtype=torch.bfloat16) + target = torch.randn(128, 16, 512, device=device, dtype=torch.bfloat16) + + losses = [] + for iteration in range(3): + optimizer.zero_grad() + with te.autocast(enabled=True, recipe=identity_recipe): + output = model(input_data) + loss = F.mse_loss(output, target) + loss.backward() + optimizer.step() + loss_val = loss.item() + assert math.isfinite(loss_val), f"Non-finite Identity loss at iter {iteration}: {loss_val}" + losses.append(loss_val) + dist_print(f"Identity iteration {iteration} completed with loss {loss_val}") + + assert losses[-1] < losses[0], f"Identity loss did not decrease: {losses}" + assert ( + _identity_param_count() == identity_count + ), "IdentityTensor params lost their quantized type after optimizer.step()" + + _check_fp8_fsdp2_allgather(model, tols={}) + + +def test_distributed_hybrid_reshard_after_forward(hybrid_recipe_name): + """FSDP2 training with hybrid params and reshard_after_forward=True. + + A single LayerNormLinear as the root module gets reshard_after_forward=True + from FSDP2. This exercises the forward-reshard-backward-reshard cycle where + split/as_strided/slice dispatch ops fire every iteration. + """ + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + device = torch.device(f"cuda:{int(os.getenv('LOCAL_RANK', '0'))}") + + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + # Keep dim-0 FSDP shards aligned to Float8 block scaling's 128-row tiles. + shard_alignment = 128 * world_size + in_features = ((512 + shard_alignment - 1) // shard_alignment) * shard_alignment + out_features = in_features * 3 + assert in_features % shard_alignment == 0 + with te.quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = te.LayerNormLinear( + in_features, + out_features, + params_dtype=torch.bfloat16, + device="meta", + ) + + custom_attrs = save_custom_attrs(model) + mesh = get_device_mesh(world_size, [world_size]) + fully_shard(model, mesh=mesh) + for module in model.modules(): + if hasattr(module, "reset_parameters"): + module.reset_parameters() + restore_custom_attrs(model, custom_attrs) + + from transformer_engine.pytorch import HybridQuantizedTensor + + def _hybrid_param_count(): + return sum( + 1 + for p in model.parameters() + if isinstance(p, DTensor) and isinstance(p._local_tensor, HybridQuantizedTensor) + ) + + hybrid_count = _hybrid_param_count() + assert hybrid_count > 0, "No HybridQuantizedTensor local tensors after sharding" + + optimizer = optim.Adam(model.parameters(), lr=1e-3) + + x = torch.randn(128, 16, in_features, device=device, dtype=torch.bfloat16) + target = torch.randn(128, 16, out_features, device=device, dtype=torch.bfloat16) + + losses = [] + for iteration in range(5): + optimizer.zero_grad() + with te.autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + loss = F.mse_loss(output, target) + loss.backward() + optimizer.step() + loss_val = loss.item() + assert math.isfinite(loss_val), f"Non-finite loss at iter {iteration}: {loss_val}" + losses.append(loss_val) + dist_print(f"Hybrid reshard_after_fwd iter {iteration}, loss {loss_val:.4f}") + + # The forward-reshard-backward-reshard cycle must still train: strict decrease. + assert all( + losses[i + 1] < losses[i] for i in range(len(losses) - 1) + ), f"Loss not strictly decreasing each step: {losses}" + + # Params must survive the split/as_strided/slice reshard dispatch ops with + # their hybrid quantized type intact. + assert ( + _hybrid_param_count() == hybrid_count + ), "HybridQuantizedTensor params lost their quantized type after optimizer.step()" + + if __name__ == "__main__": sys.exit(_train(_parse_args())) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index ee6a97ffea..6e09ed316f 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -11,6 +11,7 @@ import torch import torch.distributed as dist +from transformer_engine.common.recipe import MXFP8BlockScaling from transformer_engine.pytorch.ep import ( EpBuffer, ep_bootstrap, @@ -19,21 +20,26 @@ ep_dispatch, ep_combine, symm_mem_alloc, + release_symm_mem_pool, + is_symm_backed, _ep_combine_raw, _ep_dispatch_raw, ) - ZERO_COPY = os.environ.get("NVTE_EP_ZERO_COPY", "0") == "1" +EAGER = os.environ.get("NVTE_EP_EAGER", "0") == "1" +OVERFLOW = os.environ.get("NVTE_EP_OVERFLOW", "0") == "1" # Must come after the transformer_engine import so libtransformer_engine.so is loaded. import transformer_engine_torch as tex # noqa: F401 - NUM_LOCAL_EXPERTS = 2 -HIDDEN_DIM = 32 +# MXFP8 dispatch needs HIDDEN_DIM % 512 == 0 and TOKENS_PER_RANK % 32 == 0. Defaults +# satisfy both so the MXFP8 tests run by default; override via NVTE_EP_HIDDEN_DIM / +# NVTE_EP_TOKENS_PER_RANK. +HIDDEN_DIM = int(os.environ.get("NVTE_EP_HIDDEN_DIM", "512")) TOP_K = 2 -TOKENS_PER_RANK = 4 +TOKENS_PER_RANK = int(os.environ.get("NVTE_EP_TOKENS_PER_RANK", "32")) def _zero_copy_test_include(fn): @@ -42,6 +48,30 @@ def _zero_copy_test_include(fn): return fn +def _eager_test_include(fn): + """Mark a test to run in the eager pass; others skip there.""" + fn._eager_test_include = True + return fn + + +def _overflow_test_include(fn): + """Mark a test to run in the overflow (drop-on-overflow) pass; others skip there.""" + fn._overflow_test_include = True + return fn + + +# MXFP8 grouped dispatch needs a per-expert alignment of 128, but the EP backend caches a single +# alignment per process, so alignment=128 tests cannot share a process with the alignment=0 tests. +# They run in a dedicated pass (NVTE_EP_MXFP8_PASS=1) instead. +MXFP8_PASS = os.environ.get("NVTE_EP_MXFP8_PASS", "0") == "1" + + +def _mxfp8_align_test(fn): + """Mark a test that dispatches with alignment=128; runs only in the MXFP8 pass.""" + fn._mxfp8_align_test = True + return fn + + class _StageToSymm(torch.autograd.Function): """Identity op that stages ``src`` into a symm-mem buffer; grad passes through. Lets a test feed a symm-mem-backed, autograd-tracked tensor into ep_combine. @@ -106,6 +136,16 @@ def _make_identity_inputs(rank, ep_size, device="cuda"): ) +def _degroup_mxfp8(recv_grouped, valid_counts=None): + """Dequantize a per-expert MXFP8 GroupedTensor to a dense tensor in expert-major order. + With ``valid_counts`` keep only the first ``valid_counts[e]`` rows of each padded expert + slot; otherwise return every (padded) row.""" + parts = recv_grouped.split_into_quantized_tensors() + if valid_counts is None: + return torch.cat([p.dequantize() for p in parts], dim=0) + return torch.cat([p.dequantize()[:v] for p, v in zip(parts, valid_counts)], dim=0) + + class _Cfg: rank: int world_size: int @@ -125,6 +165,10 @@ def _make_cfg() -> _Cfg: active = min(cfg.num_experts, T * cfg.ep_size * TOP_K) overconc = cfg.num_experts // active cfg.recv_capacity_per_rank = NUM_LOCAL_EXPERTS * max(T * cfg.ep_size * TOP_K, 16) * overconc * 2 + if OVERFLOW: + # Undersize recv capacity so identity routing overflows a rank's budget; + # HT requires capacity >= max_tokens_per_rank. + cfg.recv_capacity_per_rank = TOKENS_PER_RANK cfg.device = torch.device("cuda", torch.cuda.current_device()) return cfg @@ -143,41 +187,64 @@ def setUpClass(cls): cls.ep_group, num_experts=cls.cfg.num_experts, max_tokens_per_rank=TOKENS_PER_RANK, - recv_capacity_per_rank=cls.cfg.recv_capacity_per_rank, hidden_dim=HIDDEN_DIM, + num_topk=TOP_K, + # Omit recv_capacity_per_rank to select eager mode. + recv_capacity_per_rank=None if EAGER else cls.cfg.recv_capacity_per_rank, zero_copy=ZERO_COPY, + drop_on_overflow=OVERFLOW, ) def setUp(self): + # alignment=128 MXFP8 tests run only in the dedicated MXFP8 pass; everything else skips + # there (and the MXFP8 tests skip outside it) since the backend pins one alignment/process. + is_mxfp8_align = getattr(getattr(self, self._testMethodName), "_mxfp8_align_test", False) + if MXFP8_PASS and not is_mxfp8_align: + self.skipTest("only alignment=128 MXFP8 tests run in the MXFP8 pass") + if not MXFP8_PASS and is_mxfp8_align: + self.skipTest("alignment=128 MXFP8 tests run in the dedicated MXFP8 pass") + # MXFP8 quantization requires Blackwell (SM 10.0) or newer. + if is_mxfp8_align and torch.cuda.get_device_capability() < (10, 0): + self.skipTest("MXFP8 EP tests require Blackwell (SM 10.0) or newer") # Only the zero-copy-capable tests run in the zero-copy pass. if ZERO_COPY and not getattr( getattr(self, self._testMethodName), "_zero_copy_test_include", False ): self.skipTest("not exercised in zero-copy mode") + # Only the eager-capable tests run in the eager pass. + if EAGER and not getattr(getattr(self, self._testMethodName), "_eager_test_include", False): + self.skipTest("not exercised in eager mode") + # Only the overflow-capable tests run in the overflow pass. + if OVERFLOW and not getattr( + getattr(self, self._testMethodName), "_overflow_test_include", False + ): + self.skipTest("not exercised in overflow mode") def _make_buffer( self, alignment=0, top_k=TOP_K, - dispatch_recv_tokens=None, - combine_grad_expert_out=None, + dispatch_fwd_quant_recipe=None, + combine_bwd_quant_recipe=None, ): return EpBuffer( top_k=top_k, max_tokens_per_rank=TOKENS_PER_RANK, - recv_capacity_per_rank=self.cfg.recv_capacity_per_rank, hidden_dim=HIDDEN_DIM, num_local_experts=NUM_LOCAL_EXPERTS, + recv_capacity_per_rank=None if EAGER else self.cfg.recv_capacity_per_rank, alignment=alignment, - dispatch_recv_tokens=dispatch_recv_tokens, - combine_grad_expert_out=combine_grad_expert_out, + dispatch_fwd_quant_recipe=dispatch_fwd_quant_recipe, + combine_bwd_quant_recipe=combine_bwd_quant_recipe, ) def _expert_out(self, expert_out): """Stage the combine input into symm-mem under zero-copy (combine requires it).""" if not ZERO_COPY: return expert_out - symm_buf = symm_mem_alloc(tuple(expert_out.shape), expert_out.dtype, self.ep_group) + symm_buf = symm_mem_alloc( + tuple(expert_out.shape), expert_out.dtype, self.ep_group, use_pool=True + ) return _StageToSymm.apply(expert_out, symm_buf) def _stage_grad_symm(self, x, symm_buf=None): @@ -191,12 +258,12 @@ def _stage_grad_symm(self, x, symm_buf=None): return _GradToSymm.apply(x, symm_buf) def _make_raw_recv(self, dtype=torch.bfloat16): - """Raw recv tensors + token_counts for the primitive tests.""" + """Raw recv tensors + tokens_per_expert for the primitive tests.""" rc = self.cfg.recv_capacity_per_rank return ( torch.empty(rc, HIDDEN_DIM, dtype=dtype, device=self.cfg.device), torch.empty(rc, dtype=torch.float32, device=self.cfg.device), - torch.empty(NUM_LOCAL_EXPERTS, dtype=torch.int32, device=self.cfg.device), + torch.empty(NUM_LOCAL_EXPERTS, dtype=torch.int64, device=self.cfg.device), ) @staticmethod @@ -212,17 +279,100 @@ def _moe_step(self, buffer, topk_idx, tokens, w): # Prepare + @_eager_test_include def test_primitive_prepare(self): buf = self._make_buffer() topk_idx, _toks, _w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) - token_counts = ep_prepare(buf, topk_idx) + tokens_per_expert = ep_prepare(buf, topk_idx) torch.cuda.synchronize() - self.assertEqual(token_counts.shape, (NUM_LOCAL_EXPERTS,)) - local = int(token_counts.sum().item()) + self.assertEqual(tokens_per_expert.shape, (NUM_LOCAL_EXPERTS,)) + local = int(tokens_per_expert.sum().item()) total = torch.tensor([local], dtype=torch.int64, device=self.cfg.device) dist.all_reduce(total, op=dist.ReduceOp.SUM, group=self.ep_group) self.assertEqual(int(total.item()), self.cfg.world_size * TOKENS_PER_RANK * TOP_K) + @_eager_test_include + def test_eager_recv_sizing(self): + """Eager mode sizes dispatch outputs to the exact per-step recv-token total.""" + if not EAGER: + self.skipTest("eager-only assertions") + buf = self._make_buffer() + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + recv_t, recv_w, tokens_per_expert = ep_dispatch(buf, tokens, topk_idx, w) + torch.cuda.synchronize() + # The per-step recv-token total is exposed on the buffer (int64 [1]). + self.assertEqual(buf.total_recv_tokens.dtype, torch.int64) + total = int(buf.total_recv_tokens.item()) + # recv outputs are sized to the recv total, not recv_capacity_per_rank. + self.assertEqual(recv_t.shape[0], total) + self.assertEqual(recv_w.shape[0], total) + # padded total is at least the unpadded per-expert sum and within capacity. + self.assertGreaterEqual(total, int(tokens_per_expert.sum().item())) + self.assertLessEqual(total, self.cfg.recv_capacity_per_rank) + + @_eager_test_include + def test_eager_rank_with_zero_recv_tokens(self): + """Empty recv tensors remain valid through the forward and backward pipeline.""" + if not EAGER: + self.skipTest("eager-only assertions") + buf = self._make_buffer() + _topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + # Experts [0, TOP_K) are local to rank 0, so every other rank receives + # no tokens and PyTorch gives its eager recv tensors null data pointers. + topk_idx = torch.arange(TOP_K, dtype=torch.int64, device=self.cfg.device).repeat( + TOKENS_PER_RANK, 1 + ) + tokens_p = tokens.detach().clone().requires_grad_(True) + + recv_t, recv_w, tokens_per_expert = ep_dispatch(buf, tokens_p, topk_idx, w) + recv_rows = int(buf.total_recv_tokens.item()) + self.assertEqual(recv_t.shape, (recv_rows, HIDDEN_DIM)) + self.assertEqual(recv_w.shape, (recv_rows,)) + if self.cfg.rank == 0: + self.assertGreater(recv_rows, 0) + self.assertEqual( + int(tokens_per_expert.sum().item()), + self.cfg.world_size * TOKENS_PER_RANK * TOP_K, + ) + else: + self.assertEqual(recv_rows, 0) + self.assertEqual(int(tokens_per_expert.sum().item()), 0) + self.assertEqual(recv_t.data_ptr(), 0) + self.assertEqual(recv_w.data_ptr(), 0) + + expert_out = self._weighted(recv_t, recv_w) + result = ep_combine(buf, expert_out, num_local_tokens=TOKENS_PER_RANK) + (0.5 * (result.float() ** 2).sum()).backward() + torch.cuda.synchronize() + torch.testing.assert_close(result.float(), tokens.float(), atol=5e-2, rtol=5e-2) + torch.testing.assert_close(tokens_p.grad.float(), tokens.float(), atol=5e-2, rtol=5e-2) + + @_overflow_test_include + def test_overflow_drop(self): + """drop_on_overflow: recv past capacity is dropped and dispatch continues + instead of trapping; the pre-drop recv total exceeds recv_capacity.""" + if not OVERFLOW: + self.skipTest("overflow-only assertions") + buf = self._make_buffer() + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + # Identity routing sends TOKENS_PER_RANK * TOP_K tokens to each rank, which + # overflows the deliberately undersized capacity. + expected_recv = TOKENS_PER_RANK * TOP_K + self.assertGreater(expected_recv, self.cfg.recv_capacity_per_rank) + # total_recv_tokens reports the true (pre-drop) recv total, counting the + # tokens that will be dropped; the per-expert counts exclude them and sum + # to the kept tokens (capped at recv_capacity_per_rank). + tokens_per_expert = ep_prepare(buf, topk_idx) + torch.cuda.synchronize() + self.assertEqual(int(buf.total_recv_tokens.item()), expected_recv) + self.assertEqual(int(tokens_per_expert.sum().item()), self.cfg.recv_capacity_per_rank) + # Dispatch drops overflowing tokens and completes (no trap); recv outputs + # stay capped at recv_capacity_per_rank. + recv_t, recv_w, _ = ep_dispatch(buf, tokens, topk_idx, w) + torch.cuda.synchronize() + self.assertEqual(recv_t.shape[0], self.cfg.recv_capacity_per_rank) + self.assertEqual(recv_w.shape[0], self.cfg.recv_capacity_per_rank) + # Identity round-trip via raw primitives def test_primitive_dispatch_combine_identity(self): @@ -253,10 +403,10 @@ def test_dispatch_autograd(self): ] for label, recv_tokens in cases: with self.subTest(case=label): - buf = self._make_buffer(dispatch_recv_tokens=recv_tokens) + buf = self._make_buffer() topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) tokens_p = tokens.detach().clone().requires_grad_(True) - rt, rw, _tc = ep_dispatch(buf, tokens_p, topk_idx, w) + rt, rw, _tc = ep_dispatch(buf, tokens_p, topk_idx, w, recv_tokens=recv_tokens) if recv_tokens is not None: # caller-supplied recv_tokens must be used in place self.assertEqual(rt.data_ptr(), recv_tokens.data_ptr()) rt = self._stage_grad_symm(rt) @@ -267,22 +417,83 @@ def test_dispatch_autograd(self): tokens_p.grad.float(), tokens.float() * float(TOP_K), atol=5e-2, rtol=5e-2 ) + # MXFP8 dispatch + + def _mxfp8_quantizer(self): + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + return MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False) + + def _require_mxfp8_shapes(self): + if HIDDEN_DIM % 512 != 0 or TOKENS_PER_RANK % 32 != 0: + self.skipTest( + "MXFP8 needs HIDDEN_DIM % 512 == 0 and TOKENS_PER_RANK % 32 == 0 " + "(set NVTE_EP_HIDDEN_DIM / NVTE_EP_TOKENS_PER_RANK)" + ) + + def _assert_mxfp8_matches_bf16(self, recv_mx, tokens, topk_idx, w, tc): + """Dequantized MXFP8 recv matches a bf16 dispatch of the same tokens. Both share the + alignment=128 padded expert-major layout, so compare the full prefix [0:sum(padded)].""" + ref_tokens = self._mxfp8_quantizer().quantize(tokens).dequantize() + ref_recv, _rw, _tc = ep_dispatch(self._make_buffer(alignment=128), ref_tokens, topk_idx, w) + torch.cuda.synchronize() + n = int(tc.sum()) + torch.testing.assert_close( + _degroup_mxfp8(recv_mx).float(), ref_recv.float()[:n], atol=1e-2, rtol=1e-2 + ) + + @_eager_test_include + @_zero_copy_test_include + @_mxfp8_align_test + def test_dispatch_mxfp8(self): + """MXFP8 dispatch quantizes bf16 tokens internally; recv (a per-expert GroupedTensor) + dequantized matches a bf16 dispatch of the same tokens. Under zero-copy the recv data and + scales are symm-mem backed.""" + self._require_mxfp8_shapes() + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + buf = self._make_buffer(dispatch_fwd_quant_recipe=MXFP8BlockScaling(), alignment=128) + recv_mx, _rw, tc = ep_dispatch(buf, tokens, topk_idx, w) + if ZERO_COPY: + self.assertTrue(is_symm_backed(recv_mx.rowwise_data)) + self.assertTrue(is_symm_backed(recv_mx.scale_inv)) + self._assert_mxfp8_matches_bf16(recv_mx, tokens, topk_idx, w, tc) + + @_zero_copy_test_include + @_mxfp8_align_test + def test_caller_provides_dispatch_recv_mxfp8(self): + """One caller-supplied buffer holds the recv data followed by the e8m0 scales; ep_dispatch + slices it and the returned GroupedTensor views the data and scale regions of that buffer.""" + self._require_mxfp8_shapes() + from transformer_engine.pytorch.constants import MXFP8_BLOCK_SCALING_SIZE + + rc = self.cfg.recv_capacity_per_rank + cols = HIDDEN_DIM // MXFP8_BLOCK_SCALING_SIZE + nbytes = rc * (HIDDEN_DIM + cols) # fp8 data + e8m0 scales, one byte per element + if ZERO_COPY: + recv_buf = symm_mem_alloc((nbytes,), torch.uint8, self.ep_group) + else: + recv_buf = torch.empty(nbytes, dtype=torch.uint8, device=self.cfg.device) + buf = self._make_buffer(dispatch_fwd_quant_recipe=MXFP8BlockScaling(), alignment=128) + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + recv_mx, _rw, tc = ep_dispatch(buf, tokens, topk_idx, w, recv_tokens=recv_buf) + # the returned GroupedTensor views the caller buffer's data then scale regions + self.assertEqual(recv_mx.rowwise_data.data_ptr(), recv_buf.data_ptr()) + self.assertEqual(recv_mx.scale_inv.data_ptr(), recv_buf.data_ptr() + rc * HIDDEN_DIM) + self._assert_mxfp8_matches_bf16(recv_mx, tokens, topk_idx, w, tc) + @_zero_copy_test_include def test_caller_provides_dispatch_recv_tokens(self): - """Caller-supplied recv_tokens: EpBuffer adopts it (recv_topk_weights stays - owned) and ep_dispatch returns a view of the caller's buffer.""" + """Caller-supplied recv_tokens (symm-mem-backed in zero-copy): ep_dispatch + writes into it and returns a view of the caller's buffer.""" if ZERO_COPY: rc = self.cfg.recv_capacity_per_rank rt_buf = symm_mem_alloc((rc, HIDDEN_DIM), torch.bfloat16, self.ep_group) else: rt_buf, _rw_buf, _ = self._make_raw_recv() - buf = self._make_buffer(dispatch_recv_tokens=rt_buf) - self.assertEqual(buf.recv_tokens_symm_buf.data_ptr(), rt_buf.data_ptr()) - if ZERO_COPY: # recv_topk_weights is always buffer-owned in zero-copy - self.assertIsNotNone(buf.recv_topk_weights_symm_buf) + buf = self._make_buffer() topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) tokens_p = tokens.detach().clone().requires_grad_(True) - rt, rw, _ = ep_dispatch(buf, tokens_p, topk_idx, w) + rt, rw, _ = ep_dispatch(buf, tokens_p, topk_idx, w, recv_tokens=rt_buf) self.assertEqual(rt.data_ptr(), rt_buf.data_ptr()) rt = self._stage_grad_symm(rt) rw = self._stage_grad_symm(rw) @@ -294,22 +505,124 @@ def test_caller_provides_dispatch_recv_tokens(self): @_zero_copy_test_include def test_caller_provides_grad_expert_out(self): - """Caller-supplied grad_expert_out: EpBuffer adopts it as the combine - backward grad target (symm-mem under zero-copy).""" + """Caller-supplied grad_out (symm-mem-backed in zero-copy): ep_combine's + backward scatters the expert-out grad into it.""" rc = self.cfg.recv_capacity_per_rank if ZERO_COPY: gbuf = symm_mem_alloc((rc, HIDDEN_DIM), torch.bfloat16, self.ep_group) else: gbuf = torch.empty(rc, HIDDEN_DIM, dtype=torch.bfloat16, device=self.cfg.device) - buf = self._make_buffer(combine_grad_expert_out=gbuf) - self.assertEqual(buf.grad_expert_out_symm_buf.data_ptr(), gbuf.data_ptr()) + gbuf.zero_() + buf = self._make_buffer() topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) tokens_p = tokens.detach().clone().requires_grad_(True) recv_t, recv_w, _ = ep_dispatch(buf, tokens_p, topk_idx, w) recv_t = self._stage_grad_symm(recv_t) recv_w = self._stage_grad_symm(recv_w) expert_out = self._expert_out(self._weighted(recv_t, recv_w)) - out = ep_combine(buf, expert_out) + out = ep_combine(buf, expert_out, grad_out=gbuf) + (0.5 * (out.float() ** 2).sum()).backward() + torch.cuda.synchronize() + torch.testing.assert_close(out.float(), tokens.float(), atol=5e-2, rtol=5e-2) + torch.testing.assert_close(tokens_p.grad.float(), tokens.float(), atol=5e-2, rtol=5e-2) + # the caller-owned buffer was used as the combine-bwd scatter target + self.assertGreater(gbuf.abs().sum().item(), 0.0) + + @_zero_copy_test_include + @_mxfp8_align_test + def test_combine_bwd_mxfp8_caller_grad_out(self): + """MXFP8 combine backward into a single caller buffer sliced into data + e8m0 scales: the + returned per-expert GroupedTensor views those regions and, dequantized, matches a bf16 + combine backward reference on the same routing. Under zero-copy the caller buffer and combine + input are symm-mem backed.""" + self._require_mxfp8_shapes() + from transformer_engine.pytorch.constants import MXFP8_BLOCK_SCALING_SIZE + + rc = self.cfg.recv_capacity_per_rank + cols = HIDDEN_DIM // MXFP8_BLOCK_SCALING_SIZE + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + eo_vals = ( + torch.linspace(-0.5, 0.5, rc * HIDDEN_DIM, device=self.cfg.device) + .reshape(rc, HIDDEN_DIM) + .to(torch.bfloat16) + ) + # MXFP8 combine backward writes into one caller buffer (data then e8m0 scales) + buf_mx = self._make_buffer(combine_bwd_quant_recipe=MXFP8BlockScaling(), alignment=128) + _recv, _rw, tc = ep_dispatch(buf_mx, tokens, topk_idx, w) # seeds the routing + nbytes = rc * (HIDDEN_DIM + cols) + if ZERO_COPY: + grad_buf = symm_mem_alloc((nbytes,), torch.uint8, self.ep_group) + else: + grad_buf = torch.empty(nbytes, dtype=torch.uint8, device=self.cfg.device) + src_mx = eo_vals.detach().clone().requires_grad_(True) + out_mx = ep_combine(buf_mx, self._expert_out(src_mx), grad_out=grad_buf) + (0.5 * (out_mx.float() ** 2).sum()).backward() + g_mx = src_mx.grad # per-expert GroupedTensor viewing grad_buf + self.assertEqual(g_mx.rowwise_data.data_ptr(), grad_buf.data_ptr()) + self.assertEqual(g_mx.scale_inv.data_ptr(), grad_buf.data_ptr() + rc * HIDDEN_DIM) + # bf16 reference combine backward on the same routing + buf_bf = self._make_buffer(alignment=128) + ep_dispatch(buf_bf, tokens, topk_idx, w) + src_bf = eo_vals.detach().clone().requires_grad_(True) + out_bf = ep_combine(buf_bf, self._expert_out(src_bf)) + (0.5 * (out_bf.float() ** 2).sum()).backward() + torch.cuda.synchronize() + n = int(tc.sum()) + torch.testing.assert_close( + _degroup_mxfp8(g_mx).float(), src_bf.grad.float()[:n], atol=5e-2, rtol=5e-2 + ) + + @_eager_test_include + @_zero_copy_test_include + @_mxfp8_align_test + def test_combine_bwd_mxfp8(self): + """MXFP8 combine backward with an internally allocated grad target: the returned per-expert + GroupedTensor, dequantized, matches a bf16 combine backward reference on the same routing. + Under zero-copy the combine input is symm-mem backed. + """ + self._require_mxfp8_shapes() + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + buf_mx = self._make_buffer(combine_bwd_quant_recipe=MXFP8BlockScaling(), alignment=128) + _recv, _rw, tc = ep_dispatch(buf_mx, tokens, topk_idx, w) # seeds the routing + # Combine input rows match the recv total (per-step in eager, capacity otherwise). + rows = int(buf_mx.total_recv_tokens.item()) if EAGER else self.cfg.recv_capacity_per_rank + eo_vals = ( + torch.linspace(-0.5, 0.5, rows * HIDDEN_DIM, device=self.cfg.device) + .reshape(rows, HIDDEN_DIM) + .to(torch.bfloat16) + ) + src_mx = eo_vals.detach().clone().requires_grad_(True) + out_mx = ep_combine(buf_mx, self._expert_out(src_mx)) + (0.5 * (out_mx.float() ** 2).sum()).backward() + g_mx = src_mx.grad # per-expert GroupedTensor + # bf16 reference combine backward on the same routing + buf_bf = self._make_buffer(alignment=128) + ep_dispatch(buf_bf, tokens, topk_idx, w) + src_bf = eo_vals.detach().clone().requires_grad_(True) + out_bf = ep_combine(buf_bf, self._expert_out(src_bf)) + (0.5 * (out_bf.float() ** 2).sum()).backward() + torch.cuda.synchronize() + n = int(tc.sum()) + torch.testing.assert_close( + _degroup_mxfp8(g_mx).float(), src_bf.grad.float()[:n], atol=5e-2, rtol=5e-2 + ) + + @_zero_copy_test_include + def test_zero_copy_pool_auto_alloc(self): + """Zero-copy with recv/grad left None: ep_dispatch/ep_combine allocate their IO + tensors from the symm-mem pool (is_symm_backed). This is the primary mcore + path — mcore hands no caller buffers, TE pools them on the fly.""" + if not ZERO_COPY: + self.skipTest("zero-copy pool auto-alloc only") + buf = self._make_buffer() + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + tokens_p = tokens.detach().clone().requires_grad_(True) + recv_t, recv_w, _ = ep_dispatch(buf, tokens_p, topk_idx, w) # recv_tokens=None -> pool + self.assertTrue(is_symm_backed(recv_t)) # dispatch recv came from the symm-mem pool + recv_t = self._stage_grad_symm(recv_t) + recv_w = self._stage_grad_symm(recv_w) + expert_out = self._expert_out(self._weighted(recv_t, recv_w)) + out = ep_combine(buf, expert_out) # grad_out=None -> bwd allocs the grad from the pool (0.5 * (out.float() ** 2).sum()).backward() torch.cuda.synchronize() torch.testing.assert_close(out.float(), tokens.float(), atol=5e-2, rtol=5e-2) @@ -317,6 +630,7 @@ def test_caller_provides_grad_expert_out(self): # Multi-iter stability + @_eager_test_include def test_dispatch_autograd_multiple_iterations(self): """5 fwd+bwd iters on the same EpBuffer must be bit-stable.""" buf = self._make_buffer() @@ -401,18 +715,29 @@ def _run_1f1b(self, capture): recv = [None, None, None] # Per-microbatch grad-staging buffers, symm-mem under zero-copy and - # pre-allocated so nothing is allocated/freed mid-interleave. The recv - # outputs are owned by each EpBuffer (symm-mem under zero-copy). + # pre-allocated so nothing is allocated/freed mid-interleave. recv_w = [None, None, None] rc = self.cfg.recv_capacity_per_rank if ZERO_COPY: gbuf_t = [symm_mem_alloc((rc, H), torch.bfloat16, self.ep_group) for _ in scales] gbuf_w = [symm_mem_alloc((rc,), torch.float32, self.ep_group) for _ in scales] + # Persistent symm-mem recv buffers per microbatch: leaving recv None + # pool-allocates, which is not CUDA-graph capturable. + rbuf_t = [symm_mem_alloc((rc, H), torch.bfloat16, self.ep_group) for _ in scales] + rbuf_w = [symm_mem_alloc((rc,), torch.float32, self.ep_group) for _ in scales] else: gbuf_t = gbuf_w = [None, None, None] + rbuf_t = rbuf_w = [None, None, None] def fwd(k): - rt, rw, _ = ep_dispatch(buffers[k], tokens_p[k], idx, w) + rt, rw, _ = ep_dispatch( + buffers[k], + tokens_p[k], + idx, + w, + recv_tokens=rbuf_t[k], + recv_topk_weights=rbuf_w[k], + ) recv[k] = self._stage_grad_symm(rt, gbuf_t[k]) recv_w[k] = self._stage_grad_symm(rw, gbuf_w[k]) @@ -466,6 +791,7 @@ def zero_grads(): ) @_zero_copy_test_include + @_eager_test_include def test_combine_autograd(self): """ep_combine fwd+bwd; bwd grad target is the EpBuffer symm buffer (zc) or in-flight.""" buf = self._make_buffer() @@ -504,5 +830,7 @@ def _init_distributed(): result = runner.run(suite) dist.barrier() ep_finalize() + # Deregister symm-mem windows while the comm is still valid. + release_symm_mem_pool() dist.destroy_process_group() sys.exit(0 if result.wasSuccessful() else 1) diff --git a/tests/pytorch/distributed/run_hybrid_tp_sp.py b/tests/pytorch/distributed/run_hybrid_tp_sp.py new file mode 100644 index 0000000000..b4dfb646ef --- /dev/null +++ b/tests/pytorch/distributed/run_hybrid_tp_sp.py @@ -0,0 +1,834 @@ +#!/usr/bin/python3 + +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Distributed TP/SP coverage for hybrid quantization.""" + +import argparse +import datetime +import os +import sys +from pathlib import Path + +# This file is launched directly with torchrun, so Python only adds the +# distributed test directory to sys.path. Add tests/pytorch for the shared +# hybrid recipe helpers used by both local and distributed tests. +_TEST_ROOT = str(Path(__file__).resolve().parent.parent) +if _TEST_ROOT not in sys.path: + sys.path.append(_TEST_ROOT) + +import torch +import torch.distributed as dist +from torch import nn + +import transformer_engine.pytorch as te +from transformer_engine.common import recipe as te_recipe + +from hybrid_quantization_utils import ( + hybrid_fp8_current_e5m2_grads_qfactory, + hybrid_fp8_current_identity_qfactory, + hybrid_mxfp8_identity_qfactory, + hybrid_mxfp8_qfactory, + hybrid_nvfp4_qfactory, + hybrid_tp_mxfp8_nvfp4_qfactory, + identity_qfactory, +) +from distributed.run_layer_with_overlap import _compare_tensors + +# ── Global state ───────────────────────────────────────────────────── + +SEQ_LEN = 32 +BATCH_SIZE = 32 +HIDDEN_SIZE = 128 +FFN_HIDDEN_SIZE = 128 +NR_HEADS = 4 + +WORLD_RANK = None +WORLD_SIZE = None +NCCL_WORLD = None +QUANTIZATION = None + +LOSS_FN = nn.MSELoss() + + +def hybrid_recipe(): + """Return a fresh CustomRecipe for the selected test recipe.""" + if QUANTIZATION == "hybrid_fp8": + return te_recipe.CustomRecipe(qfactory=hybrid_fp8_current_e5m2_grads_qfactory) + if QUANTIZATION == "hybrid_mxfp8": + return te_recipe.CustomRecipe(qfactory=hybrid_mxfp8_qfactory) + if QUANTIZATION == "hybrid_fp8_identity": + return te_recipe.CustomRecipe(qfactory=hybrid_fp8_current_identity_qfactory) + if QUANTIZATION == "hybrid_mxfp8_identity": + return te_recipe.CustomRecipe(qfactory=hybrid_mxfp8_identity_qfactory) + if QUANTIZATION == "identity": + return te_recipe.CustomRecipe(qfactory=identity_qfactory) + if QUANTIZATION == "hybrid_nvfp4": + return te_recipe.CustomRecipe(qfactory=hybrid_nvfp4_qfactory) + if QUANTIZATION == "hybrid_mxfp8_nvfp4": + return te_recipe.CustomRecipe(qfactory=hybrid_tp_mxfp8_nvfp4_qfactory) + raise ValueError(f"Unknown hybrid QUANTIZATION={QUANTIZATION!r}") + + +# ── Tolerances ─────────────────────────────────────────────────────── +# +# Mostly upstream ``run_numerics.py`` tolerances. NVFP4 needs a slightly +# larger atol for the measured column+SP Linear output. + + +def _get_tolerances(): + if QUANTIZATION == "identity": + # BF16 TP reductions accumulate in a different order. + return {"rtol": 1.6e-2, "atol": 1.0e-5} + if QUANTIZATION in ("hybrid_fp8", "hybrid_fp8_identity"): + # Loose because of sequence parallel & amax reduction (fp8_cs). + return {"rtol": 0.4, "atol": 0.25} + if QUANTIZATION in ("hybrid_mxfp8", "hybrid_mxfp8_identity"): + return {"rtol": 0.125, "atol": 0.0625} + if QUANTIZATION == "hybrid_nvfp4": + # Measured column+SP Linear output max abs is ~0.144. + return {"rtol": 0.125, "atol": 0.15} + if QUANTIZATION == "hybrid_mxfp8_nvfp4": + # Backward GEMMs run in NVFP4 -> inherit the (looser) NVFP4 bounds. + return {"rtol": 0.125, "atol": 0.12} + raise ValueError(f"No tolerances for QUANTIZATION={QUANTIZATION!r}") + + +# ── Distributed helpers ────────────────────────────────────────────── + + +def dist_print(msg, src=None, error=False): + stream = sys.stderr if error else sys.stdout + if WORLD_RANK == (0 if src is None else src): + stream.write(f"[rank{WORLD_RANK}] {msg}\n") + stream.flush() + + +def _gather(tensor, dim=0): + """All-gather with ``run_numerics.py`` gradient scaling.""" + + class HalfGradient(torch.autograd.Function): + @staticmethod + def forward(ctx, inp): + return inp + + @staticmethod + def backward(ctx, grad_output): + return grad_output / WORLD_SIZE + + tensor = HalfGradient.apply(tensor) + gathered = torch.distributed.nn.functional.all_gather(tensor, group=NCCL_WORLD) + return torch.cat(gathered, dim=dim) + + +def _copy_params(model_distributed, model_single): + """Copy single-node parameters into the local TP shard.""" + for dp, sp in zip(model_distributed.parameters(), model_single.parameters()): + with torch.no_grad(): + to_copy = sp + for dim, _ in enumerate(dp.shape): + if dp.shape[dim] != sp.shape[dim]: + start = WORLD_RANK * dp.shape[dim] + end = (WORLD_RANK + 1) * dp.shape[dim] + indices = [slice(None)] * max(min(dim, len(dp.shape) - 1), 0) + indices.append(slice(start, end)) + if dim < len(dp.shape) - 1: + indices.append(slice(None)) + to_copy = sp[tuple(indices)] + dp.copy_(to_copy) + + +def _match_param_sizes(dist_param, single_param): + indices = [slice(None)] * len(single_param.shape) + for i in range(len(dist_param.shape)): + if dist_param.shape[i] != single_param.shape[i]: + start = WORLD_RANK * dist_param.shape[i] + end = (WORLD_RANK + 1) * dist_param.shape[i] + indices[i] = slice(start, end) + return single_param[tuple(indices)] + + +def _check_outputs(output_single, output_dist, label="outputs", *, tolerances=None): + failed = torch.tensor([0], dtype=torch.uint8, device="cuda") + f, info = _compare_tensors( + label, output_dist, output_single, **(tolerances or _get_tolerances()) + ) + if f: + dist_print(info, src=WORLD_RANK, error=True) + failed[0] = int(f) + dist.all_reduce(failed, dist.ReduceOp.MAX, NCCL_WORLD) + assert not bool(failed.item()), f"{label}: numerical check failed on at least one rank" + + +def _collective_assert(condition, label): + """Raise on every rank only after all ranks report structural status.""" + failed = torch.tensor([not condition], dtype=torch.uint8, device="cuda") + if not condition: + dist_print(label, src=WORLD_RANK, error=True) + dist.all_reduce(failed, dist.ReduceOp.MAX, NCCL_WORLD) + assert not bool(failed.item()), f"{label}: failed on at least one rank" + + +def _check_gradients( + model_dist, model_single, *, bitwise=False, reduce_replicated=False, tolerances=None +): + """Compare every parameter gradient, including presence and parameter count. + + Sequence-parallel replicated parameters accumulate a partial gradient on + each rank. ``reduce_replicated`` reconstructs their full reference gradient + without mutating the gradient held by the model. + """ + dist_params = list(model_dist.named_parameters()) + single_params = list(model_single.named_parameters()) + local_counts = torch.tensor( + [len(dist_params), len(single_params)], dtype=torch.int64, device="cuda" + ) + min_counts = local_counts.clone() + max_counts = local_counts.clone() + dist.all_reduce(min_counts, dist.ReduceOp.MIN, NCCL_WORLD) + dist.all_reduce(max_counts, dist.ReduceOp.MAX, NCCL_WORLD) + _collective_assert( + torch.equal(min_counts, max_counts) and len(dist_params) == len(single_params), + "parameter count mismatch: " + f"local distributed={len(dist_params)}, local single={len(single_params)}, " + f"global min={min_counts.tolist()}, global max={max_counts.tolist()}", + ) + for i, ((name, pd), (single_name, ps)) in enumerate(zip(dist_params, single_params)): + _collective_assert( + name == single_name, + f"parameter {i} name mismatch: {name!r} != {single_name!r}", + ) + local_presence = torch.tensor( + [pd.grad is not None, ps.grad is not None], dtype=torch.uint8, device="cuda" + ) + min_presence = local_presence.clone() + max_presence = local_presence.clone() + dist.all_reduce(min_presence, dist.ReduceOp.MIN, NCCL_WORLD) + dist.all_reduce(max_presence, dist.ReduceOp.MAX, NCCL_WORLD) + _collective_assert( + torch.equal(min_presence, max_presence) + and local_presence[0].item() == local_presence[1].item(), + f"grad[{i}].{name}: gradient presence differs locally or across ranks: " + f"local={local_presence.tolist()}, global min={min_presence.tolist()}, " + f"global max={max_presence.tolist()}", + ) + if pd.grad is None: + continue + pd_grad = pd.grad + local_reduce = torch.tensor( + [reduce_replicated and pd_grad.shape == ps.grad.shape], + dtype=torch.uint8, + device="cuda", + ) + min_reduce = local_reduce.clone() + max_reduce = local_reduce.clone() + dist.all_reduce(min_reduce, dist.ReduceOp.MIN, NCCL_WORLD) + dist.all_reduce(max_reduce, dist.ReduceOp.MAX, NCCL_WORLD) + _collective_assert( + torch.equal(min_reduce, max_reduce), + f"grad[{i}].{name}: replicated-gradient reduction branch differs across ranks", + ) + if local_reduce.item(): + pd_grad = pd_grad.detach().clone() + dist.all_reduce(pd_grad, dist.ReduceOp.SUM, NCCL_WORLD) + ps_grad = _match_param_sizes(pd_grad, ps.grad) + label = f"grad[{i}].{name}" + if bitwise: + _check_bitwise(pd_grad, ps_grad, label) + else: + _check_outputs(ps_grad, pd_grad, label=label, tolerances=tolerances) + + +def _check_input_gradient(inp_dist, inp_single, label, *, bitwise=False): + """Compare the local TP/SP input-gradient shard with its full reference.""" + local_presence = torch.tensor( + [inp_dist.grad is not None, inp_single.grad is not None], + dtype=torch.uint8, + device="cuda", + ) + min_presence = local_presence.clone() + max_presence = local_presence.clone() + dist.all_reduce(min_presence, dist.ReduceOp.MIN, NCCL_WORLD) + dist.all_reduce(max_presence, dist.ReduceOp.MAX, NCCL_WORLD) + _collective_assert( + torch.equal(min_presence, max_presence) + and local_presence[0].item() == local_presence[1].item(), + f"{label}: input-gradient presence differs locally or across ranks: " + f"local={local_presence.tolist()}, global min={min_presence.tolist()}, " + f"global max={max_presence.tolist()}", + ) + _collective_assert( + inp_dist.grad is not None, + f"{label}: both input gradients are unexpectedly absent", + ) + expected = _match_param_sizes(inp_dist.grad, inp_single.grad) + if bitwise: + _check_bitwise(inp_dist.grad, expected, label) + else: + _check_outputs(expected, inp_dist.grad, label=label) + + +def _apply_models(model_single, model_dist, inp_single, inp_dist, **kwargs): + """Run both models with fresh CustomRecipe instances.""" + inp_single.requires_grad_() + inp_dist.requires_grad_() + with te.autocast(enabled=True, recipe=hybrid_recipe()): + out_single = model_single(inp_single, **kwargs) + with te.autocast(enabled=True, recipe=hybrid_recipe()): + out_dist = model_dist(inp_dist, **kwargs) + return out_single, out_dist + + +def _loss_backward(out_single, out_dist): + target = torch.randn_like(out_single) + LOSS_FN(out_single, target).backward() + LOSS_FN(out_dist, target).backward() + + +# ── Test 1: te.Linear TP (column + row) × SP (on/off) ──────────────── + + +def _test_linear(parallel_mode, sequence_parallel, params_dtype=torch.bfloat16, amax_stress=False): + dist_print( + f"linear: parallel_mode={parallel_mode} sequence_parallel={sequence_parallel}" + f" dtype={params_dtype} amax_stress={amax_stress}" + ) + + torch.manual_seed(12345) + torch.cuda.manual_seed(12345) + + model_single = te.Linear(HIDDEN_SIZE, HIDDEN_SIZE, params_dtype=params_dtype).cuda() + model_dist = te.Linear( + HIDDEN_SIZE, + HIDDEN_SIZE, + tp_size=WORLD_SIZE, + tp_group=NCCL_WORLD, + parallel_mode=parallel_mode, + sequence_parallel=sequence_parallel, + params_dtype=params_dtype, + ).cuda() + + _copy_params(model_dist, model_single) + + # Match run_numerics._test_linear input layouts. + inp_single = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + if parallel_mode == "row": + split = HIDDEN_SIZE // WORLD_SIZE + inp_dist = inp_single[:, WORLD_RANK * split : (WORLD_RANK + 1) * split].clone() + elif parallel_mode == "column": + if sequence_parallel: + # SP column: input is sharded along batch/sequence dim 0. + inp_single = torch.empty((WORLD_SIZE * BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + inp_dist = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + if amax_stress and WORLD_RANK == WORLD_SIZE - 1: + # One-rank outlier before SP gather. + inp_dist[-1, -1] = 1.0e3 + inp_single = _gather(inp_dist, dim=0).detach() + else: + inp_dist = inp_single.clone() + else: + raise ValueError(parallel_mode) + + out_single, out_dist = _apply_models(model_single, model_dist, inp_single, inp_dist) + + # For column-parallel: output is split along feature dim 1; gather. + # For row-parallel + SP: output is split along seq dim 0; gather. + if parallel_mode == "column" or (sequence_parallel and parallel_mode == "row"): + gather_dim = 1 if parallel_mode == "column" else 0 + out_dist = _gather(out_dist, dim=gather_dim) + + _loss_backward(out_single, out_dist) + _check_outputs( + out_single, + out_dist, + label=f"linear[{parallel_mode},sp={sequence_parallel},amax_stress={amax_stress}]", + ) + + _check_input_gradient( + inp_dist, + inp_single, + label=f"linear[{parallel_mode},sp={sequence_parallel},amax_stress={amax_stress}] dgrad", + ) + _check_gradients(model_dist, model_single, reduce_replicated=sequence_parallel) + + +def test_linear(): + for parallel_mode in ["column", "row"]: + for sequence_parallel in [False, True]: + _test_linear(parallel_mode, sequence_parallel) + # Current-scaling amax stress: one rank owns an outlier before SP gather. + if QUANTIZATION in ("hybrid_fp8", "hybrid_fp8_identity"): + _test_linear("column", True, amax_stress=True) + + +# ── Test 1b: te.Linear hybrid-vs-vanilla bitwise operand equivalence ─ + + +def vanilla_recipe(): + """Built-in recipe for same-format hybrid-vs-vanilla checks.""" + if QUANTIZATION == "hybrid_fp8": + return te_recipe.Float8CurrentScaling() + if QUANTIZATION == "hybrid_mxfp8": + return te_recipe.MXFP8BlockScaling() + if QUANTIZATION == "hybrid_nvfp4": + return te_recipe.NVFP4BlockScaling() + if QUANTIZATION == "identity": + return None + raise ValueError(f"No vanilla recipe for QUANTIZATION={QUANTIZATION!r}") + + +def _backward_not_bitwise_comparable(): + """NVFP4 backward consumes SR RNG differently in hybrid vs vanilla.""" + return QUANTIZATION == "hybrid_nvfp4" + + +def _check_bitwise(actual, expected, label): + """Assert bitwise equality (rtol=0, atol=0), all-reduced across ranks.""" + failed = torch.tensor([0], dtype=torch.uint8, device="cuda") + try: + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + except AssertionError as exc: + dist_print(f"{label}: {exc}", src=WORLD_RANK, error=True) + failed[0] = 1 + dist.all_reduce(failed, dist.ReduceOp.MAX, NCCL_WORLD) + assert not bool(failed.item()), f"{label}: not bitwise-identical on at least one rank" + + +def _test_linear_vs_vanilla(parallel_mode, sequence_parallel, params_dtype=torch.bfloat16): + """Same-topology Linear check: forward bitwise, backward where comparable.""" + dist_print( + f"linear_vs_vanilla: parallel_mode={parallel_mode} sequence_parallel={sequence_parallel}" + ) + + def run(recipe): + # Fresh model per recipe (re-seeded for identical weights): TE caches a + # quantized weight workspace on the module, so reusing one model would + # let the first recipe's cached weight contaminate the second. + torch.manual_seed(12345) + torch.cuda.manual_seed(12345) + model = te.Linear( + HIDDEN_SIZE, + HIDDEN_SIZE, + tp_size=WORLD_SIZE, + tp_group=NCCL_WORLD, + parallel_mode=parallel_mode, + sequence_parallel=sequence_parallel, + params_dtype=params_dtype, + ).cuda() + + torch.manual_seed(34567) + torch.cuda.manual_seed(34567) + inp = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + if parallel_mode == "row": + split = HIDDEN_SIZE // WORLD_SIZE + inp = inp[:, WORLD_RANK * split : (WORLD_RANK + 1) * split].clone() + inp.requires_grad_() + + with te.autocast(enabled=recipe is not None, recipe=recipe): + out = model(inp) + # Fixed, recipe-independent target so both backward graphs match. + torch.manual_seed(54321) + torch.cuda.manual_seed(54321) + target = torch.randn_like(out) + LOSS_FN(out, target).backward() + weight_grads = [p.grad.detach().clone() for p in model.parameters() if p.grad is not None] + return out.detach().clone(), inp.grad.detach().clone(), weight_grads + + out_h, dinp_h, wgrads_h = run(hybrid_recipe()) + out_v, dinp_v, wgrads_v = run(vanilla_recipe()) + + tag = f"linear_vs_vanilla[{parallel_mode},sp={sequence_parallel}]" + + # Same-topology fprop operands should match bitwise. + _check_bitwise(out_h, out_v, f"{tag} forward") + + # NVFP4 backward consumes columnwise stochastic-rounding RNG differently. + if not _backward_not_bitwise_comparable(): + _check_bitwise(dinp_h, dinp_v, f"{tag} dgrad") + assert len(wgrads_h) == len(wgrads_v), f"{tag}: weight-grad count mismatch" + for i, (gh, gv) in enumerate(zip(wgrads_h, wgrads_v)): + _check_bitwise(gh, gv, f"{tag} wgrad[{i}]") + + +def test_linear_vs_vanilla(): + # These recipes have no same-format vanilla bitwise target. + if QUANTIZATION in ( + "hybrid_mxfp8_nvfp4", + "hybrid_fp8_identity", + "hybrid_mxfp8_identity", + ): + dist_print("linear_vs_vanilla: skipped for hybrid without a vanilla equivalent") + return + for parallel_mode in ["column", "row"]: + for sequence_parallel in [False, True]: + _test_linear_vs_vanilla(parallel_mode, sequence_parallel) + + +def _same_format_parity_supported(): + return QUANTIZATION in ("hybrid_fp8", "hybrid_mxfp8", "identity") + + +def _check_same_topology_parity( + out_h, dinp_h, model_h, out_v, dinp_v, model_v, tag, *, tolerances=None +): + check = ( + _check_bitwise + if tolerances is None + else lambda a, e, label: _check_outputs(e, a, label=label, tolerances=tolerances) + ) + check(out_h, out_v, f"{tag} forward") + check(dinp_h, dinp_v, f"{tag} dgrad") + _check_gradients(model_h, model_v, bitwise=tolerances is None, tolerances=tolerances) + + +def _test_layernorm_linear_vs_vanilla(sequence_parallel, params_dtype=torch.bfloat16): + if not _same_format_parity_supported(): + dist_print("layernorm_linear_vs_vanilla: skipped for recipe without vanilla equivalent") + return + dist_print(f"layernorm_linear_vs_vanilla: sequence_parallel={sequence_parallel}") + + def run(recipe_obj): + torch.manual_seed(23456) + torch.cuda.manual_seed(23456) + model = te.LayerNormLinear( + HIDDEN_SIZE, + HIDDEN_SIZE, + tp_size=WORLD_SIZE, + tp_group=NCCL_WORLD, + parallel_mode="column", + sequence_parallel=sequence_parallel, + params_dtype=params_dtype, + ).cuda() + torch.manual_seed(45670) + torch.cuda.manual_seed(45670) + inp = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + inp.requires_grad_() + with te.autocast(enabled=recipe_obj is not None, recipe=recipe_obj): + out = model(inp) + torch.manual_seed(45671) + torch.cuda.manual_seed(45671) + LOSS_FN(out, torch.randn_like(out)).backward() + return model, out.detach().clone(), inp.grad.detach().clone() + + model_h, out_h, dinp_h = run(hybrid_recipe()) + model_v, out_v, dinp_v = run(vanilla_recipe()) + _check_same_topology_parity( + out_h, + dinp_h, + model_h, + out_v, + dinp_v, + model_v, + f"layernorm_linear_vs_vanilla[sp={sequence_parallel}]", + ) + + +def test_layernorm_linear_vs_vanilla(): + for sequence_parallel in [False, True]: + _test_layernorm_linear_vs_vanilla(sequence_parallel) + + +def _test_layernorm_mlp_vs_vanilla(sequence_parallel, params_dtype=torch.bfloat16): + if not _same_format_parity_supported(): + dist_print("layernorm_mlp_vs_vanilla: skipped for recipe without vanilla equivalent") + return + dist_print(f"layernorm_mlp_vs_vanilla: sequence_parallel={sequence_parallel}") + + def run(recipe_obj): + torch.manual_seed(45678) + torch.cuda.manual_seed(45678) + model = te.LayerNormMLP( + HIDDEN_SIZE, + FFN_HIDDEN_SIZE, + tp_size=WORLD_SIZE, + tp_group=NCCL_WORLD, + set_parallel_mode=True, + sequence_parallel=sequence_parallel, + params_dtype=params_dtype, + ).cuda() + torch.manual_seed(56780) + torch.cuda.manual_seed(56780) + inp = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + inp.requires_grad_() + with te.autocast(enabled=recipe_obj is not None, recipe=recipe_obj): + out = model(inp) + torch.manual_seed(56781) + torch.cuda.manual_seed(56781) + LOSS_FN(out, torch.randn_like(out)).backward() + return model, out.detach().clone(), inp.grad.detach().clone() + + model_h, out_h, dinp_h = run(hybrid_recipe()) + model_v, out_v, dinp_v = run(vanilla_recipe()) + _check_same_topology_parity( + out_h, + dinp_h, + model_h, + out_v, + dinp_v, + model_v, + f"layernorm_mlp_vs_vanilla[sp={sequence_parallel}]", + # Hybrid CustomRecipe disables quantized-norm fusion while the native + # recipe uses it. Both paths consume identical quantized GEMM operands, + # but the BF16 norm boundary can differ by one rounding step. + tolerances=(None if QUANTIZATION == "identity" else {"rtol": 2**-7, "atol": 2**-10}), + ) + + +def test_layernorm_mlp_vs_vanilla(): + for sequence_parallel in [False, True]: + _test_layernorm_mlp_vs_vanilla(sequence_parallel) + + +# ── Test 2: te.LayerNormLinear column + SP ────────────────────────── + + +def _test_layernorm_linear(sequence_parallel, params_dtype=torch.bfloat16): + """Column-parallel LayerNormLinear with optional SP.""" + dist_print(f"layernorm_linear: parallel_mode=column sequence_parallel={sequence_parallel}") + + torch.manual_seed(23456) + torch.cuda.manual_seed(23456) + + model_single = te.LayerNormLinear(HIDDEN_SIZE, HIDDEN_SIZE, params_dtype=params_dtype).cuda() + model_dist = te.LayerNormLinear( + HIDDEN_SIZE, + HIDDEN_SIZE, + tp_size=WORLD_SIZE, + tp_group=NCCL_WORLD, + parallel_mode="column", + sequence_parallel=sequence_parallel, + params_dtype=params_dtype, + ).cuda() + + _copy_params(model_dist, model_single) + + if sequence_parallel: + inp_dist = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + inp_single = _gather(inp_dist, dim=0).detach() + else: + inp_single = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + inp_dist = inp_single.clone() + + out_single, out_dist = _apply_models(model_single, model_dist, inp_single, inp_dist) + + # Column-parallel output: gather along dim 1. + out_dist = _gather(out_dist, dim=1) + + _loss_backward(out_single, out_dist) + _check_outputs(out_single, out_dist, label=f"layernorm_linear[sp={sequence_parallel}]") + + _check_input_gradient( + inp_dist, inp_single, label=f"layernorm_linear[sp={sequence_parallel}] dgrad" + ) + _check_gradients(model_dist, model_single, reduce_replicated=sequence_parallel) + + +def test_layernorm_linear(): + for sequence_parallel in [False, True]: + _test_layernorm_linear(sequence_parallel) + + +# ── Test 3: te.LayerNormMLP + TP + SP ─────────────────────────────── + + +def _test_layernorm_mlp(sequence_parallel, params_dtype=torch.bfloat16): + """LayerNormMLP with set_parallel_mode=True and optional SP.""" + dist_print(f"layernorm_mlp: parallel_mode=set sequence_parallel={sequence_parallel}") + + torch.manual_seed(45678) + torch.cuda.manual_seed(45678) + + model_single = te.LayerNormMLP(HIDDEN_SIZE, FFN_HIDDEN_SIZE, params_dtype=params_dtype).cuda() + model_dist = te.LayerNormMLP( + HIDDEN_SIZE, + FFN_HIDDEN_SIZE, + tp_size=WORLD_SIZE, + tp_group=NCCL_WORLD, + set_parallel_mode=True, + sequence_parallel=sequence_parallel, + params_dtype=params_dtype, + ).cuda() + + _copy_params(model_dist, model_single) + + if sequence_parallel: + inp_dist = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + inp_single = _gather(inp_dist, dim=0).detach() + else: + inp_single = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + inp_dist = inp_single.clone() + + out_single, out_dist = _apply_models(model_single, model_dist, inp_single, inp_dist) + + # Row-parallel FC2 output is in the full hidden space; with SP it is + # reduce-scattered along the token dim 0, so gather it back. + if sequence_parallel: + out_dist = _gather(out_dist, dim=0) + + _loss_backward(out_single, out_dist) + _check_outputs(out_single, out_dist, label=f"layernorm_mlp[sp={sequence_parallel}]") + + _check_input_gradient( + inp_dist, inp_single, label=f"layernorm_mlp[sp={sequence_parallel}] dgrad" + ) + _check_gradients(model_dist, model_single, reduce_replicated=sequence_parallel) + + +def test_layernorm_mlp(): + for sequence_parallel in [False, True]: + _test_layernorm_mlp(sequence_parallel) + + +# ── Test 4: te.TransformerLayer + TP + SP ─────────────────────────── + + +def _test_transformer_layer(sequence_parallel, params_dtype=torch.bfloat16): + """TransformerLayer integration with TP and optional SP.""" + dist_print(f"transformer_layer: parallel_mode=set sequence_parallel={sequence_parallel}") + + torch.manual_seed(34567) + torch.cuda.manual_seed(34567) + + model_single = te.TransformerLayer( + HIDDEN_SIZE, + FFN_HIDDEN_SIZE, + NR_HEADS, + attention_dropout=0.0, + hidden_dropout=0.0, + fuse_qkv_params=True, + params_dtype=params_dtype, + ).cuda() + model_dist = te.TransformerLayer( + HIDDEN_SIZE, + FFN_HIDDEN_SIZE, + NR_HEADS, + tp_size=WORLD_SIZE, + tp_group=NCCL_WORLD, + set_parallel_mode=True, + sequence_parallel=sequence_parallel, + seq_length=WORLD_SIZE * SEQ_LEN if sequence_parallel else None, + attention_dropout=0.0, + hidden_dropout=0.0, + fuse_qkv_params=True, + params_dtype=params_dtype, + ).cuda() + + _copy_params(model_dist, model_single) + + inp_single = ( + torch.randn((WORLD_SIZE * SEQ_LEN, BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + ) + if sequence_parallel: + inp_dist = inp_single[WORLD_RANK * SEQ_LEN : (WORLD_RANK + 1) * SEQ_LEN, :, :].contiguous() + else: + inp_dist = inp_single.clone() + + out_single, out_dist = _apply_models(model_single, model_dist, inp_single, inp_dist) + + if sequence_parallel: + out_dist = _gather(out_dist, dim=0) + + _loss_backward(out_single, out_dist) + _check_outputs(out_single, out_dist, label=f"transformer_layer[sp={sequence_parallel}]") + + _check_input_gradient( + inp_dist, inp_single, label=f"transformer_layer[sp={sequence_parallel}] dgrad" + ) + _check_gradients(model_dist, model_single, reduce_replicated=sequence_parallel) + + +def test_transformer_layer(): + for sequence_parallel in [False, True]: + _test_transformer_layer(sequence_parallel) + + +# ── Driver ─────────────────────────────────────────────────────────── + + +def main(argv=None): + global WORLD_RANK, WORLD_SIZE, NCCL_WORLD, QUANTIZATION + + WORLD_RANK = int(os.getenv("RANK", "0")) + WORLD_SIZE = int(os.getenv("WORLD_SIZE", "1")) + LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0")) + LOCAL_SIZE = int(os.getenv("LOCAL_WORLD_SIZE", "1")) + + assert WORLD_SIZE == LOCAL_SIZE, "This test is single-node only" + assert LOCAL_SIZE <= torch.cuda.device_count() + + torch.cuda.set_device(LOCAL_RANK) + dist.init_process_group( + backend="nccl", + rank=WORLD_RANK, + world_size=WORLD_SIZE, + timeout=datetime.timedelta(seconds=60), + init_method="env://", + device_id=torch.device(f"cuda:{LOCAL_RANK}"), + ) + NCCL_WORLD = dist.new_group(backend="nccl") + + parser = argparse.ArgumentParser() + parser.add_argument( + "--quantization", + type=str, + required=True, + choices=[ + "hybrid_fp8", + "hybrid_mxfp8", + "hybrid_fp8_identity", + "hybrid_mxfp8_identity", + "identity", + "hybrid_nvfp4", + "hybrid_mxfp8_nvfp4", + ], + ) + parser.add_argument( + "--test", + type=str, + nargs="+", + default=["all"], + choices=[ + "all", + "linear", + "linear_vs_vanilla", + "layernorm_linear_vs_vanilla", + "layernorm_mlp_vs_vanilla", + "layernorm_linear", + "layernorm_mlp", + "transformer_layer", + ], + help="Run one or more named tests in the same distributed process group", + ) + args = parser.parse_args(argv) + QUANTIZATION = args.quantization + + test_map = { + "linear": test_linear, + "linear_vs_vanilla": test_linear_vs_vanilla, + "layernorm_linear_vs_vanilla": test_layernorm_linear_vs_vanilla, + "layernorm_mlp_vs_vanilla": test_layernorm_mlp_vs_vanilla, + "layernorm_linear": test_layernorm_linear, + "layernorm_mlp": test_layernorm_mlp, + "transformer_layer": test_transformer_layer, + } + if "all" in args.test: + if len(args.test) != 1: + parser.error("--test all cannot be combined with named tests") + tests_to_run = list(test_map.values()) + else: + tests_to_run = [test_map[name] for name in args.test] + + for test_fn in tests_to_run: + dist_print(f"=== Starting {test_fn.__name__} ===") + test_fn() + dist.barrier() + dist_print(f"=== Passed {test_fn.__name__} ===") + + dist.destroy_process_group() + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/pytorch/distributed/run_layer_with_overlap.py b/tests/pytorch/distributed/run_layer_with_overlap.py index 61573f3837..8eea0dda10 100644 --- a/tests/pytorch/distributed/run_layer_with_overlap.py +++ b/tests/pytorch/distributed/run_layer_with_overlap.py @@ -19,6 +19,7 @@ import torch import torch.distributed as dist +from torch.distributed.elastic.multiprocessing.errors import record import transformer_engine.pytorch as te from transformer_engine.common.recipe import ( @@ -363,6 +364,7 @@ def _compare_tensors(name, test, ref, rtol, atol): return numerics_failed, numerics_info +@record def _train(opts): if "OMPI_COMM_WORLD_SIZE" in os.environ: # Execution with `mpirun -np N` diff --git a/tests/pytorch/distributed/run_newton_schulz.py b/tests/pytorch/distributed/run_newton_schulz.py index 712d83bd1c..bd061949ad 100644 --- a/tests/pytorch/distributed/run_newton_schulz.py +++ b/tests/pytorch/distributed/run_newton_schulz.py @@ -8,7 +8,9 @@ """ import argparse -import sys +import itertools +import os +from typing import Optional import torch import torch.distributed as dist @@ -18,9 +20,14 @@ CusolverMpCtx, get_coefficients, newton_schulz, + newton_schulz_tp, ) +DTYPES = ("float32", "bfloat16") +COEFFICIENT_CONFIGS = ((5, "quintic"), (8, "polar_express")) + + def newton_schulz_reference( in_x: torch.Tensor, coefficients: list[tuple[float, float, float]] ) -> torch.Tensor: @@ -32,34 +39,46 @@ def newton_schulz_reference( return x -@record -def main(): - parser = argparse.ArgumentParser(description="Newton-Schulz distributed test") - parser.add_argument( - "--check", type=str, default="orthogonality", choices=["orthogonality", "reference"] - ) - parser.add_argument("--dtype", type=str, default="float32", choices=["float32", "bfloat16"]) - parser.add_argument("--matrix-rows", type=int, default=256) - parser.add_argument("--matrix-cols", type=int, default=None) - parser.add_argument("--num-iterations", type=int, default=5) - parser.add_argument("--coeff-type", type=str, default="quintic") - parser.add_argument("--atol", type=float, default=1e-2) - parser.add_argument("--rtol", type=float, default=1e-2) - args = parser.parse_args() +def _dtype_from_name(dtype: str) -> torch.dtype: + if dtype == "float32": + return torch.float32 + if dtype == "bfloat16": + return torch.bfloat16 + raise ValueError(f"Unsupported dtype: {dtype}") - dist.init_process_group(backend="nccl") - rank = dist.get_rank() - world_size = dist.get_world_size() - torch.cuda.set_device(rank) - dtype = torch.float32 if args.dtype == "float32" else torch.bfloat16 - m = args.matrix_rows - n = args.matrix_cols if args.matrix_cols is not None else args.matrix_rows - coefficients = get_coefficients(args.num_iterations, args.coeff_type) +def _test_tolerances(dtype: str, check: str, world_size: int) -> tuple[float, float]: + if dtype == "bfloat16": + return (5e-2, 5e-2) + if check == "orthogonality" and world_size == 1: + return (2e-2, 2e-2) + return (1e-2, 1e-2) + - # Ensure the distributed column dimension is divisible by world_size. - assert n % world_size == 0, f"Matrix columns {n} must be divisible by world_size {world_size}" +def _shape_scale(world_size: int) -> int: + return 4 if world_size == 1 else world_size + +def _orthogonality_shapes(world_size: int) -> list[tuple[int, int]]: + scale = _shape_scale(world_size) + return [ + (scale * 64, scale * 64), + (scale * 64, scale * 96), + (scale * 96, scale * 64), + ] + + +def _reference_shapes(world_size: int) -> list[tuple[int, int]]: + scale = _shape_scale(world_size) + return [(scale * 64, scale * 64)] + + +def _make_matrix( + m: int, + n: int, + dtype: torch.dtype, + rank: int, +) -> torch.Tensor: # Create a random matrix on rank 0 with singular values in (0, 1), # which keeps the Newton-Schulz iterations in the convergence regime. if rank == 0: @@ -72,56 +91,201 @@ def main(): torch.randn(n, k, device="cuda", dtype=torch.float32), mode="reduced" ) singular_values = torch.rand(k, device="cuda", dtype=torch.float32) * 0.8 + 0.1 - A = U @ torch.diag(singular_values) @ V.T - A = A.to(dtype) + matrix = U @ torch.diag(singular_values) @ V.T + matrix = matrix.to(dtype) else: - A = torch.empty(m, n, device="cuda", dtype=dtype) + matrix = torch.empty(m, n, device="cuda", dtype=dtype) - # Broadcast the full matrix to all ranks - dist.broadcast(A, src=0) + dist.broadcast(matrix, src=0) + return matrix - # Scatter columns to each rank - local_cols = n // world_size - x_local = A[:, rank * local_cols : (rank + 1) * local_cols].contiguous() - ctx = CusolverMpCtx(dist.group.WORLD) - try: - newton_schulz(x_local, ctx, args.num_iterations, coefficients=coefficients) - finally: - ctx.destroy() +def _run_case( + *, + ctx: CusolverMpCtx, + check: str, + dtype_name: str, + matrix_shape: tuple[int, int], + num_iterations: int, + coeff_type: str, + api: str = "base", + partition_dim: Optional[int] = 1, + tp_mode: str = "distributed", +) -> None: + rank = ctx.rank + world_size = ctx.nranks + dtype = _dtype_from_name(dtype_name) + m, n = matrix_shape + coefficients = get_coefficients(num_iterations, coeff_type) + atol, rtol = _test_tolerances(dtype_name, check, world_size) + + if api == "tp" and partition_dim is None: + # Replicated inputs are sharded along the larger dimension for cuSolverMp. + assert ( + max(m, n) % world_size == 0 + ), f"Matrix dimension {max(m, n)} must be divisible by world_size {world_size}" + elif api == "base" or partition_dim == 1: + # Ensure the distributed column dimension is divisible by world_size. + assert ( + n % world_size == 0 + ), f"Matrix columns {n} must be divisible by world_size {world_size}" + else: + assert m % world_size == 0, f"Matrix rows {m} must be divisible by world_size {world_size}" + + A = _make_matrix(m, n, dtype, rank) + + # Replicate the full tensor or scatter it along the API's partition dimension. + if api == "tp" and partition_dim is None: + x_local = A.clone() + gather_dim = None + elif api == "tp" and partition_dim == 0: + local_rows = m // world_size + x_local = A[rank * local_rows : (rank + 1) * local_rows, :].contiguous().clone() + gather_dim = 0 + else: + local_cols = n // world_size + x_local = A[:, rank * local_cols : (rank + 1) * local_cols].contiguous().clone() + gather_dim = 1 - # Gather results - gathered = [torch.empty_like(x_local) for _ in range(world_size)] - dist.all_gather(gathered, x_local) - X = torch.cat(gathered, dim=1) + if api == "tp": + newton_schulz_tp( + x_local, + ctx, + num_iterations, + coefficients=coefficients, + partition_dim=partition_dim, + tp_mode=tp_mode, + ) + else: + newton_schulz(x_local, ctx, num_iterations, coefficients=coefficients) + + # Reconstruct the full result unless the TP API already returned a replicated tensor. + if gather_dim is None: + X = x_local + else: + gathered = [torch.empty_like(x_local) for _ in range(world_size)] + dist.all_gather(gathered, x_local) + X = torch.cat(gathered, dim=gather_dim) # Check: the resulting matrix should be orthogonal, or match a local reference. - if rank == 0: - if args.check == "orthogonality": - if m <= n: - gram = X @ X.t() - expected = torch.eye(m, device=gram.device, dtype=gram.dtype) - max_diff = (gram - expected).abs().max().item() - print(f"Max |X @ X.t() - I|: {max_diff:.6e}", flush=True) - else: - gram = X.t() @ X - expected = torch.eye(n, device=gram.device, dtype=gram.dtype) - max_diff = (gram - expected).abs().max().item() - print(f"Max |X.t() @ X - I|: {max_diff:.6e}", flush=True) - passed = torch.allclose(gram, expected, atol=args.atol, rtol=args.rtol) + if check == "orthogonality": + if m <= n: + gram = X @ X.t() + expected = torch.eye(m, device=gram.device, dtype=gram.dtype) + label = "X @ X.t() - I" else: - reference = newton_schulz_reference(A.float(), coefficients).to(dtype) - max_diff = (X - reference).abs().max().item() - print(f"Max |distributed - reference|: {max_diff:.6e}", flush=True) - passed = torch.allclose(X, reference, atol=args.atol, rtol=args.rtol) + gram = X.t() @ X + expected = torch.eye(n, device=gram.device, dtype=gram.dtype) + label = "X.t() @ X - I" + max_diff = (gram - expected).abs().max().item() + passed = torch.allclose(gram, expected, atol=atol, rtol=rtol) + elif check == "reference": + reference = newton_schulz_reference(A.float(), coefficients).to(dtype) + max_diff = (X - reference).abs().max().item() + label = "distributed - reference" + passed = torch.allclose(X, reference, atol=atol, rtol=rtol) + else: + raise ValueError(f"Unsupported check: {check}") - if passed: - print("NUMERICAL CHECK PASSED", flush=True) - else: - print("NUMERICAL CHECK FAILED", flush=True, file=sys.stderr) - sys.exit(1) + if rank == 0: + print(f"Max |{label}|: {max_diff:.6e}", flush=True) + + if not passed: + raise AssertionError( + "Newton-Schulz case failed: " + f"check={check}, dtype={dtype_name}, matrix_shape={matrix_shape}, " + f"num_iterations={num_iterations}, coeff_type={coeff_type}, api={api}, " + f"partition_dim={partition_dim}, tp_mode={tp_mode}, max_diff={max_diff:.6e}" + ) + + +def run_all_tests(ctx: CusolverMpCtx) -> None: + """Run all distributed Newton-Schulz checks in one torchrun invocation.""" + rank = ctx.rank + world_size = ctx.nranks + + for config in itertools.product( + DTYPES, + _orthogonality_shapes(world_size), + COEFFICIENT_CONFIGS, + ): + dtype_name, matrix_shape, (num_iterations, coeff_type) = config + if rank == 0: + print(f"Running orthogonality check with {config=}", flush=True) + _run_case( + ctx=ctx, + check="orthogonality", + dtype_name=dtype_name, + matrix_shape=matrix_shape, + num_iterations=num_iterations, + coeff_type=coeff_type, + ) + + for config in itertools.product( + DTYPES, + _reference_shapes(world_size), + COEFFICIENT_CONFIGS, + ): + dtype_name, matrix_shape, (num_iterations, coeff_type) = config + if rank == 0: + print(f"Running reference check with {config=}", flush=True) + _run_case( + ctx=ctx, + check="reference", + dtype_name=dtype_name, + matrix_shape=matrix_shape, + num_iterations=num_iterations, + coeff_type=coeff_type, + ) + + for partition_dim, tp_mode in itertools.product((0, 1), ("duplicated", "distributed")): + config = (partition_dim, tp_mode) + if rank == 0: + print(f"Running TP API reference check with {config=}", flush=True) + _run_case( + ctx=ctx, + check="reference", + dtype_name="float32", + matrix_shape=_reference_shapes(world_size)[0], + num_iterations=5, + coeff_type="quintic", + api="tp", + partition_dim=partition_dim, + tp_mode=tp_mode, + ) + + if rank == 0: + print("Running TP API reference check with replicated input", flush=True) + _run_case( + ctx=ctx, + check="reference", + dtype_name="float32", + matrix_shape=_reference_shapes(world_size)[0], + num_iterations=5, + coeff_type="quintic", + api="tp", + partition_dim=None, + ) + + if rank == 0: + print("NUMERICAL CHECK PASSED", flush=True) + + +@record +def main(): + parser = argparse.ArgumentParser(description="Newton-Schulz distributed test") + parser.parse_args() + + dist.init_process_group(backend="nccl") + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) - dist.destroy_process_group() + ctx = CusolverMpCtx(dist.group.WORLD) + try: + run_all_tests(ctx) + finally: + ctx.destroy() + dist.destroy_process_group() if __name__ == "__main__": diff --git a/tests/pytorch/distributed/run_numerics_exact.py b/tests/pytorch/distributed/run_numerics_exact.py index aa0825338d..a5d6db7074 100644 --- a/tests/pytorch/distributed/run_numerics_exact.py +++ b/tests/pytorch/distributed/run_numerics_exact.py @@ -23,8 +23,8 @@ ) from transformer_engine.pytorch import NVFP4Quantizer from transformer_engine.pytorch.constants import NVFP4_BLOCK_SCALING_SIZE -from transformer_engine.pytorch.custom_recipes import quantization_ref_nvfp4 -from transformer_engine.pytorch.custom_recipes import utils +from transformer_engine.pytorch.custom_recipes import reference_nvfp4 +from transformer_engine.pytorch.custom_recipes import reference_utils # Executed as a script, so sibling imports rely on the interpreter putting this file's # directory on sys.path -- which safe-path mode (PYTHONSAFEPATH, python -P) disables. @@ -63,7 +63,7 @@ def get_nvfp4_quantizer_factory(): Mirrors the canonical "branch on what we care about, default fall-through" pattern from - ``transformer_engine.pytorch.custom_recipes.quantization_recipes_base``; + ``transformer_engine.pytorch.custom_recipes.quantizer_factories``; every slot gets a real :class:`NVFP4QuantizerRef` (``CustomRecipeState`` rejects ``None`` returns). @@ -78,14 +78,14 @@ def factory(role): and role.tensor_type == "weight" ) if is_weight: - return quantization_ref_nvfp4.NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + return reference_nvfp4.NVFP4QuantizerRef( + dtype=reference_utils.Fp4Formats.E2M1, quant_tile_shape=(16, 16), pow_2_scales=False, with_rht=False, ) - return quantization_ref_nvfp4.NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + return reference_nvfp4.NVFP4QuantizerRef( + dtype=reference_utils.Fp4Formats.E2M1, quant_tile_shape=(1, 16), pow_2_scales=False, with_rht=True, diff --git a/tests/pytorch/distributed/run_test_ep.sh b/tests/pytorch/distributed/run_test_ep.sh index 68b691f787..5788c2ac9f 100755 --- a/tests/pytorch/distributed/run_test_ep.sh +++ b/tests/pytorch/distributed/run_test_ep.sh @@ -9,8 +9,6 @@ set -uo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" -export PYTHONPATH="${TE_REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) if [ "${DETECTED_GPUS}" -lt 4 ]; then @@ -47,11 +45,15 @@ RET=0 run_pass() { local label="$1" local zc="$2" + local eager="${3:-0}" + local overflow="${4:-0}" + local mxfp8="${5:-0}" local log="stdout_ep_${label}.txt" echo "=== Running ${SCRIPT} [${label}] on ${NUM_RANKS} GPUs (timeout=${TEST_TIMEOUT_S}s) ===" # setsid + kill-after so SIGKILL takes down the whole process group, not just torchrun. - NVTE_EP_ZERO_COPY="${zc}" setsid timeout --foreground --kill-after=10 --signal=TERM \ - "${TEST_TIMEOUT_S}" \ + NVTE_EP_ZERO_COPY="${zc}" NVTE_EP_EAGER="${eager}" NVTE_EP_OVERFLOW="${overflow}" \ + NVTE_EP_MXFP8_PASS="${mxfp8}" \ + setsid timeout --foreground --kill-after=10 --signal=TERM "${TEST_TIMEOUT_S}" \ torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_RANKS}" \ "${SCRIPT}" 2>&1 | tee "${log}" local rc=${PIPESTATUS[0]} @@ -70,5 +72,13 @@ run_pass() { run_pass "default" 0 run_pass "zero_copy" 1 +run_pass "eager" 0 1 +run_pass "overflow" 0 0 1 +# MXFP8 grouped dispatch pins the per-expert alignment to 128, which the backend caches +# process-wide, so its tests get their own passes (normal + zero-copy + eager IO). mxfp8 is the +# 5th arg; eager is the 3rd. +run_pass "mxfp8" 0 0 0 1 +run_pass "mxfp8_zero_copy" 1 0 0 1 +run_pass "mxfp8_eager" 0 1 0 1 exit $RET diff --git a/tests/pytorch/distributed/test_comm_gemm_overlap.py b/tests/pytorch/distributed/test_comm_gemm_overlap.py index 902f19e3d4..78eab92975 100644 --- a/tests/pytorch/distributed/test_comm_gemm_overlap.py +++ b/tests/pytorch/distributed/test_comm_gemm_overlap.py @@ -45,6 +45,24 @@ if tex.ubuf_built_with_mpi(): LAUNCH_CMD = ["mpirun", "-np", str(NUM_PROCS), "--oversubscribe", "--quiet", "python3"] +OUTPUT_TAIL_CHARS = 4000 + + +def _assert_subprocess_succeeded(result): + if ( + result.returncode != 0 + or "NUMERICAL CHECK FAILED" in result.stderr + or "NUMERICAL CHECK PASSED" not in result.stdout + ): + raise AssertionError( + f"Distributed test exited with return code {result.returncode}" + f"\n--- stdout (last {OUTPUT_TAIL_CHARS} characters) ---\n" + f"{result.stdout[-OUTPUT_TAIL_CHARS:]}" + f"\n--- stderr (last {OUTPUT_TAIL_CHARS} characters) ---\n" + f"{result.stderr[-OUTPUT_TAIL_CHARS:]}" + ) + + # Fall back on CUDA IPC if the platform does not support CUDA multicast if not tex.device_supports_multicast(): os.environ["UB_SKIPMC"] = "1" @@ -102,13 +120,8 @@ def _run_gemm_with_overlap( ) test_cmd.append("--use-cublasmp") - result = subprocess.run(test_cmd, env=os.environ, capture_output=True, check=False) - if ( - result.returncode != 0 - or "NUMERICAL CHECK FAILED" in result.stderr.decode() - or "NUMERICAL CHECK PASSED" not in result.stdout.decode() - ): - raise AssertionError(result.stderr.decode()) + result = subprocess.run(test_cmd, env=os.environ, capture_output=True, text=True, check=False) + _assert_subprocess_succeeded(result) def _run_layer_with_overlap( @@ -164,33 +177,23 @@ def _run_layer_with_overlap( pytest.skip("cuBLASMp comm+GEMM overlap does not yet support MXFP8 (block scaling).") test_cmd.append("--use-cublasmp") - os.environ["PYTORCH_JIT"] = "0" - os.environ["NVTE_TORCH_COMPILE"] = "0" - os.environ["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "0" + test_env = os.environ.copy() + test_env["PYTORCH_JIT"] = "0" + test_env["NVTE_TORCH_COMPILE"] = "0" + test_env["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "0" if te.get_device_compute_capability() <= (8, 0): # We've experienced numerical discrepancies in Flash Attention # backward when running with Userbuffers on A100s. This does # not show up in more recent GPUs. - os.environ["NVTE_FLASH_ATTN"] = "0" + test_env["NVTE_FLASH_ATTN"] = "0" elif fp8: # Fused attention is causing non-deterministic FP8 failures on H100s even with # NVTE_ALLOW_NONDETERMINISTIC_ALGO=0, so disable it entirely for this test. - os.environ["NVTE_FUSED_ATTN"] = "0" - - result = subprocess.run(test_cmd, env=os.environ, capture_output=True, check=False) + test_env["NVTE_FUSED_ATTN"] = "0" - os.unsetenv("PYTORCH_JIT") - os.unsetenv("NVTE_TORCH_COMPILE") - os.unsetenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO") - os.unsetenv("NVTE_FLASH_ATTN") - os.unsetenv("NVTE_FUSED_ATTN") + result = subprocess.run(test_cmd, env=test_env, capture_output=True, text=True, check=False) - if ( - result.returncode != 0 - or "NUMERICAL CHECK FAILED" in result.stderr.decode() - or "NUMERICAL CHECK PASSED" not in result.stdout.decode() - ): - raise AssertionError(result.stderr.decode()) + _assert_subprocess_succeeded(result) @pytest.mark.parametrize("use_cublasmp", (False,) if IS_HIP_EXTENSION else (False, True)) diff --git a/tests/pytorch/distributed/test_ep.py b/tests/pytorch/distributed/test_ep.py index 81eef9a3c1..3978c6908c 100644 --- a/tests/pytorch/distributed/test_ep.py +++ b/tests/pytorch/distributed/test_ep.py @@ -15,17 +15,29 @@ LAUNCHER = TEST_ROOT / "run_test_ep.sh" +def _count_launcher_passes() -> int: + # Count run_pass invocations so the outer timeout scales as passes are added. + n = 0 + for line in LAUNCHER.read_text().splitlines(): + s = line.strip() + if s.startswith("run_pass ") or s.startswith("run_pass\t"): + n += 1 + return max(n, 1) + + @pytest.mark.skipif(torch.cuda.device_count() < 4, reason="EP requires >= 4 GPUs") def test_multi_process_ep(): """Launch the EP unit-test suite across all visible GPUs. - Short timeout so a hang on any rank surfaces fast rather than burning CI time. + Per-pass timeout stays short so a hang on any rank surfaces fast; the outer + pytest budget scales with the number of passes the launcher runs. """ - timeout_s = int(os.environ.get("NVTE_TEST_EP_TIMEOUT_S", "180")) + per_pass_s = int(os.environ.get("NVTE_TEST_EP_TIMEOUT_S", "180")) + outer_s = per_pass_s * _count_launcher_passes() + 60 proc = subprocess.run( ["bash", str(LAUNCHER)], - env={**os.environ, "KEEP_EP_LOGS": "1", "TEST_TIMEOUT_S": str(timeout_s)}, - timeout=timeout_s + 30, + env={**os.environ, "KEEP_EP_LOGS": "1", "TEST_TIMEOUT_S": str(per_pass_s)}, + timeout=outer_s, check=False, ) assert proc.returncode == 0, f"EP test suite failed (rc={proc.returncode})" diff --git a/tests/pytorch/distributed/test_fusible_ops.py b/tests/pytorch/distributed/test_fusible_ops.py index 8335e69b50..ab8e5b84a4 100644 --- a/tests/pytorch/distributed/test_fusible_ops.py +++ b/tests/pytorch/distributed/test_fusible_ops.py @@ -35,7 +35,8 @@ # Import utility functions _current_file = pathlib.Path(__file__).resolve() -sys.path.append(str(_current_file.parent.parent)) +# Prepend so installed packages with a top-level utils module cannot shadow the test helpers. +sys.path = [str(_current_file.parent.parent)] + sys.path from utils import dtype_tols, make_recipe, quantization_tols diff --git a/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py b/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py index eb43ba7e75..07dffebf5f 100644 --- a/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py +++ b/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py @@ -37,7 +37,8 @@ # Import utility functions _current_file = pathlib.Path(__file__).resolve() -sys.path.append(str(_current_file.parent.parent)) +# Prepend so installed packages with a top-level utils module cannot shadow the test helpers. +sys.path = [str(_current_file.parent.parent)] + sys.path from utils import dtype_tols, make_recipe, run_distributed, str_to_dtype # Check if FP8 is supported diff --git a/tests/pytorch/distributed/test_hybrid_tp_sp.py b/tests/pytorch/distributed/test_hybrid_tp_sp.py new file mode 100644 index 0000000000..6030e9d32f --- /dev/null +++ b/tests/pytorch/distributed/test_hybrid_tp_sp.py @@ -0,0 +1,146 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Pytest launcher for hybrid TP/SP distributed tests.""" + +import os +import subprocess +from pathlib import Path + +import pytest +import torch +import transformer_engine.pytorch as te +from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported + +if torch.cuda.device_count() < 2: + pytest.skip( + "Distributed TP/SP tests need at least 2 GPUs.", + allow_module_level=True, + ) + +fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) +mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) +nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) + +TEST_ROOT = Path(__file__).parent.resolve() +NUM_PROCS = min(2, torch.cuda.device_count()) +LAUNCH_CMD = ["torchrun", f"--nproc_per_node={NUM_PROCS}"] + +xfail_hopper_columnwise_per_tensor_fp8 = pytest.mark.xfail( + condition=not is_non_tn_fp8_gemm_supported(), + strict=True, + reason=( + "Hopper does not yet support columnwise-only per-tensor FP8 quantization; " + "tracked by NVIDIA/TransformerEngine#3158" + ), +) + + +def _run_tests(quantization: str, tests: tuple[str, ...] = ("all",)): + """Run related cases under one torchrun/process-group startup.""" + script = TEST_ROOT / "run_hybrid_tp_sp.py" + cmd = LAUNCH_CMD + [ + str(script), + "--quantization", + quantization, + "--test", + *tests, + ] + result = subprocess.run(cmd, env=os.environ, check=False) + assert result.returncode == 0, ( + f"run_hybrid_tp_sp.py (quantization={quantization}, tests={list(tests)})" + f" exited with code {result.returncode}" + ) + + +# ────────────────────────────────────────────────────────────────────── +# Hybrid FP8 current scaling +# ────────────────────────────────────────────────────────────────────── +# Exercises TP amax reduction and SP gather paths. + + +_SAME_FORMAT_TESTS = ( + "linear", + "linear_vs_vanilla", + "layernorm_linear_vs_vanilla", + "layernorm_mlp_vs_vanilla", + "layernorm_linear", + "layernorm_mlp", + "transformer_layer", +) + + +@pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") +@xfail_hopper_columnwise_per_tensor_fp8 +def test_hybrid_fp8(): + """Hybrid FP8 TP/SP coverage and same-topology vanilla parity.""" + _run_tests("hybrid_fp8", _SAME_FORMAT_TESTS) + + +@pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") +def test_hybrid_fp8_identity_linear(): + """Linear TP/SP coverage for FP8 forward plus Identity backward.""" + _run_tests("hybrid_fp8_identity", ("linear",)) + + +@pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") +def test_identity_all_modules(): + """All-Identity TP/SP end-to-end coverage for every supported TE module.""" + _run_tests("identity") + + +# ────────────────────────────────────────────────────────────────────── +# Hybrid MXFP8 +# ────────────────────────────────────────────────────────────────────── +# Covers per-block scale layout through TP shards. + + +@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") +def test_hybrid_mxfp8(): + _run_tests("hybrid_mxfp8", _SAME_FORMAT_TESTS) + + +@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") +def test_hybrid_mxfp8_identity_linear(): + """Linear TP/SP coverage for MXFP8 forward plus Identity backward.""" + _run_tests("hybrid_mxfp8_identity", ("linear",)) + + +# ────────────────────────────────────────────────────────────────────── +# Hybrid NVFP4 +# ────────────────────────────────────────────────────────────────────── +# Same-format NVFP4 coverage with base role-wise settings. + + +@pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}") +def test_hybrid_nvfp4(): + _run_tests( + "hybrid_nvfp4", + ( + "linear", + "linear_vs_vanilla", + "layernorm_linear", + "layernorm_mlp", + "transformer_layer", + ), + ) + + +# ────────────────────────────────────────────────────────────────────── +# Cross-format hybrid: MXFP8 rowwise + NVFP4 columnwise +# ────────────────────────────────────────────────────────────────────── +# No single vanilla recipe exists for bitwise comparison. + +_cross_format_available = mxfp8_available and nvfp4_available +_reason_for_no_cross_format = reason_for_no_mxfp8 if not mxfp8_available else reason_for_no_nvfp4 + + +@pytest.mark.skipif( + not _cross_format_available, reason=f"MXFP8+NVFP4: {_reason_for_no_cross_format}" +) +def test_hybrid_mxfp8_nvfp4(): + _run_tests( + "hybrid_mxfp8_nvfp4", + ("linear", "layernorm_linear", "layernorm_mlp", "transformer_layer"), + ) diff --git a/tests/pytorch/distributed/test_newton_schulz.py b/tests/pytorch/distributed/test_newton_schulz.py index c7fa0ff11d..be413955a8 100644 --- a/tests/pytorch/distributed/test_newton_schulz.py +++ b/tests/pytorch/distributed/test_newton_schulz.py @@ -7,14 +7,16 @@ import mmap import os import subprocess +import sys from pathlib import Path import pytest import torch from transformer_engine.common import _get_shared_object_file -if torch.cuda.device_count() < 2: - pytest.skip("Newton-Schulz tests require at least 2 GPUs.", allow_module_level=True) +NUM_PROCS = torch.cuda.device_count() +if NUM_PROCS < 1: + pytest.skip("Newton-Schulz tests require at least 1 GPU.", allow_module_level=True) def _built_with_cusolvermp() -> bool: @@ -35,32 +37,24 @@ def _built_with_cusolvermp() -> bool: ) TEST_ROOT = Path(__file__).parent.resolve() -NUM_PROCS = torch.cuda.device_count() -LAUNCH_CMD = ["torchrun", f"--nproc_per_node={NUM_PROCS}"] -ORTHOGONALITY_SHAPES = [ - (NUM_PROCS * 64, NUM_PROCS * 64), - (NUM_PROCS * 64, NUM_PROCS * 96), - (NUM_PROCS * 96, NUM_PROCS * 64), -] -REFERENCE_SHAPES = [(NUM_PROCS * 64, NUM_PROCS * 64)] -def _run_test(dtype, matrix_shape, num_iterations, coeff_type, check): - rows, cols = matrix_shape +def _run_worker(num_procs: int) -> None: test_path = TEST_ROOT / "run_newton_schulz.py" - test_cmd = LAUNCH_CMD + [ + test_cmd = [ + sys.executable, + "-m", + "torch.distributed.run", + f"--nproc_per_node={num_procs}", str(test_path), - f"--check={check}", - f"--dtype={dtype}", - f"--matrix-rows={rows}", - f"--matrix-cols={cols}", - f"--num-iterations={num_iterations}", - f"--coeff-type={coeff_type}", ] - if dtype == "bfloat16": - test_cmd += ["--atol=5e-2", "--rtol=5e-2"] - - result = subprocess.run(test_cmd, env=os.environ, capture_output=True, check=False, timeout=300) + result = subprocess.run( + test_cmd, + env=os.environ, + capture_output=True, + check=False, + timeout=1200, + ) if ( result.returncode != 0 or "NUMERICAL CHECK FAILED" in result.stderr.decode() @@ -73,17 +67,14 @@ def _run_test(dtype, matrix_shape, num_iterations, coeff_type, check): ) -@pytest.mark.parametrize("dtype", ["float32", "bfloat16"]) -@pytest.mark.parametrize("matrix_shape", ORTHOGONALITY_SHAPES) -@pytest.mark.parametrize("num_iterations,coeff_type", [(5, "quintic"), (8, "polar_express")]) -def test_orthogonality(dtype, matrix_shape, num_iterations, coeff_type): - """Test distributed Newton-Schulz orthogonality.""" - _run_test(dtype, matrix_shape, num_iterations, coeff_type, "orthogonality") +def test_newton_schulz_single_gpu(): + """Test cuSolverMp Newton-Schulz with a single-rank GPU grid.""" + _run_worker(1) -@pytest.mark.parametrize("dtype", ["float32", "bfloat16"]) -@pytest.mark.parametrize("matrix_shape", REFERENCE_SHAPES) -@pytest.mark.parametrize("num_iterations,coeff_type", [(5, "quintic"), (8, "polar_express")]) -def test_against_reference(dtype, matrix_shape, num_iterations, coeff_type): - """Test distributed Newton-Schulz against a local reference implementation.""" - _run_test(dtype, matrix_shape, num_iterations, coeff_type, "reference") +@pytest.mark.skipif( + NUM_PROCS < 2, reason="Distributed Newton-Schulz tests require at least 2 GPUs." +) +def test_newton_schulz_distributed(): + """Launch one parallel job that runs all multi-GPU Newton-Schulz checks.""" + _run_worker(NUM_PROCS) diff --git a/tests/pytorch/distributed/test_sanity.py b/tests/pytorch/distributed/test_sanity.py index 2e7a63e0a2..1b8bf5890f 100644 --- a/tests/pytorch/distributed/test_sanity.py +++ b/tests/pytorch/distributed/test_sanity.py @@ -19,7 +19,8 @@ from transformer_engine.common import recipe _current_file = pathlib.Path(__file__).resolve() -sys.path.append(str(_current_file.parent.parent)) +# Prepend so installed packages with a top-level utils module cannot shadow the test helpers. +sys.path = [str(_current_file.parent.parent)] + sys.path from utils import ModelConfig model_configs = { diff --git a/tests/pytorch/distributed/test_torch_fsdp2.py b/tests/pytorch/distributed/test_torch_fsdp2.py index 71a8dcb216..3a45a5453a 100644 --- a/tests/pytorch/distributed/test_torch_fsdp2.py +++ b/tests/pytorch/distributed/test_torch_fsdp2.py @@ -47,6 +47,8 @@ def test_fsdp2_model_tests(): "-v", "-s", "--tb=short", + "-k", + "not hybrid", ], valid_returncodes=(0, 5), env=os.environ, @@ -74,7 +76,7 @@ def test_fsdp2_fused_adam_tests(): # The following 2 tests need to be run in sequence, # as they depend on each other. "-k", - "not dcp_resharding_save and not dcp_resharding_load", + "not hybrid and not dcp_resharding_save and not dcp_resharding_load", ], valid_returncodes=(0, 5), env=os.environ, @@ -170,6 +172,57 @@ def test_fsdp2_fused_adam_dcp_resharding(recipe): assert result.returncode == 0, f"DCP resharding load phase failed: {result.returncode}" +@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") +@pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") +def test_fsdp2_hybrid_fused_adam_tests(): + """FSDP2 FusedAdam tests with hybrid quantized params (parametrized by hybrid recipe).""" + test_path = _FSDP2_DIR / "run_fsdp2_fused_adam.py" + nproc = min(NUM_PROCS, 2) + run_distributed( + [ + "torchrun", + f"--nproc_per_node={nproc}", + "--local-ranks-filter=0", + "-m", + "pytest", + str(test_path), + "-v", + "-s", + "--tb=short", + "-k", + "hybrid", + ], + valid_returncodes=(0, 5), + env=os.environ, + timeout=600, + ) + + +@pytest.mark.skipif(NUM_PROCS % 2 != 0, reason="Requires even number of GPUs") +@pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") +def test_fsdp2_hybrid_model_tests(): + """FSDP2 model tests with hybrid quantized params (parametrized by hybrid recipe).""" + test_path = _FSDP2_DIR / "run_fsdp2_model.py" + run_distributed( + [ + "torchrun", + f"--nproc_per_node={NUM_PROCS}", + "--local-ranks-filter=0", + "-m", + "pytest", + str(test_path), + "-v", + "-s", + "--tb=short", + "-k", + "hybrid", + ], + valid_returncodes=(0, 5), + env=os.environ, + timeout=600, + ) + + def test_dummy() -> None: """Dummy test diff --git a/tests/pytorch/hybrid_quantization_utils.py b/tests/pytorch/hybrid_quantization_utils.py new file mode 100644 index 0000000000..8ef4a7f421 --- /dev/null +++ b/tests/pytorch/hybrid_quantization_utils.py @@ -0,0 +1,377 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Shared explicit quantizer factories for hybrid quantization tests. + +These are intentionally ordinary module-level functions rather than a public +factory abstraction. Keeping them module-level makes ``CustomRecipe`` objects +that reference them picklable in distributed checkpoint tests. +""" + +import torch + +import transformer_engine.pytorch as te +from transformer_engine.common import recipe +from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + current_scaling_factory, + float8_block_scaling_factory, + mxfp8_factory, + nvfp4_factory, +) + +_LINEAR_MODULE_TYPES = ("linear", "grouped_linear") +_FORWARD_TENSOR_TYPES = ("input", "weight", "output") +_GRAD_TENSOR_TYPES = ("grad_output", "grad_input") + + +def _is_linear_role(role): + return role is not None and role.module_type in _LINEAR_MODULE_TYPES + + +def _make_fp8_current(*, fp8_dtype=te.DType.kFloat8E4M3): + return te.Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device="cuda") + + +def _make_mxfp8(*, fp8_dtype=te.DType.kFloat8E4M3): + return te.MXFP8Quantizer(fp8_dtype=fp8_dtype) + + +def nvfp4_linear_mxfp8_dpa_test_factory(role): + """Test-local reference for the experimental NVFP4 Linear + MXFP8 DPA example.""" + is_dpa = role is not None and role.module_type == "dpa" + is_dpa_boundary = ( + role is not None + and not role.module_type + and ("dpa_output" in role.name or "dpa_grad_input" in role.name) + ) + + if is_dpa or is_dpa_boundary: + is_bwd_role = (is_dpa and role.tensor_type in ("do", "dp", "dqkv")) or ( + is_dpa_boundary and "dpa_grad_input" in role.name + ) + fp8_dtype = te.DType.kFloat8E5M2 if is_bwd_role else te.DType.kFloat8E4M3 + return te.MXFP8Quantizer(fp8_dtype=fp8_dtype) + + return nvfp4_factory(role) + + +def fp8_e4m3_factory(): + """Construct the default E4M3 current-scaling test quantizer.""" + return te.Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") + + +def fp8_e5m2_factory(): + """Construct the default E5M2 current-scaling test quantizer.""" + return te.Float8CurrentScalingQuantizer(te.DType.kFloat8E5M2, device="cuda") + + +def mxfp8_e4m3_factory(): + """Construct the default E4M3 MXFP8 test quantizer.""" + return te.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) + + +def make_fp8_quantizer(*, rowwise=True, columnwise=True): + """Construct the standard current-scaling FP8 test quantizer.""" + return te.Float8CurrentScalingQuantizer( + fp8_dtype=te.DType.kFloat8E4M3, + device="cuda", + rowwise=rowwise, + columnwise=columnwise, + ) + + +def make_nvfp4_quantizer(*, rowwise=True, columnwise=True): + """Construct the standard NVFP4 test quantizer.""" + return te.NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=rowwise, + columnwise=columnwise, + ) + + +def make_hybrid_quantizer_fp8_row_fp4_col(): + """Construct a hybrid quantizer with FP8 rowwise and NVFP4 columnwise.""" + return te.HybridQuantizer( + rowwise_quantizer=make_fp8_quantizer(), + columnwise_quantizer=make_nvfp4_quantizer(), + ) + + +def as_data_tensor_tuple(storage): + """Return a storage's raw buffers as a tuple without copying them.""" + if storage is None: + return () + tensors = storage.get_data_tensors() + return tensors if isinstance(tensors, tuple) else (tensors,) + + +def snapshot_storage_tensor_metadata(storage, *, clone=False): + """Capture every tensor-valued concrete-storage metadata field.""" + if storage is None: + return None + tensor_metadata = {} + for name, value in storage.get_metadata().items(): + if isinstance(value, torch.Tensor) or value is None: + snapshot_value = value.detach().clone() if clone and value is not None else value + if ( + value is not None + and getattr(storage, "_is_2D_scaled", False) + and name in ("rowwise_scale_inv", "columnwise_scale_inv") + ): + # Float8Block 2D scales pad one tile dimension. Kernels do not + # initialize or consume that padding, so canonicalize it. + snapshot_value = value.detach().clone() + m, n = storage._fsdp_logical_mn() + block_len = storage._FSDP_BLOCK_LEN + m_tiles = (m + block_len - 1) // block_len + n_tiles = (n + block_len - 1) // block_len + if name == "rowwise_scale_inv": + snapshot_value[:, n_tiles:] = 0 + else: + snapshot_value[:, m_tiles:] = 0 + tensor_metadata[name] = snapshot_value + return { + "storage_type": type(storage).__name__, + "tensor_metadata": tensor_metadata, + } + + +def assert_nested_state_exact(actual, expected, *, path="state"): + """Recursively compare nested state, using zero tolerance for tensors.""" + if isinstance(expected, torch.Tensor): + assert isinstance(actual, torch.Tensor), f"{path}: expected Tensor, got {type(actual)}" + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0, msg=path) + return + if isinstance(expected, dict): + assert isinstance(actual, dict), f"{path}: expected dict, got {type(actual)}" + assert actual.keys() == expected.keys(), f"{path}: dictionary keys differ" + for key in expected: + assert_nested_state_exact(actual[key], expected[key], path=f"{path}.{key}") + return + if isinstance(expected, (list, tuple)): + assert isinstance( + actual, type(expected) + ), f"{path}: expected {type(expected).__name__}, got {type(actual).__name__}" + assert len(actual) == len(expected), f"{path}: sequence lengths differ" + for index, (actual_item, expected_item) in enumerate(zip(actual, expected)): + assert_nested_state_exact( + actual_item, + expected_item, + path=f"{path}[{index}]", + ) + return + assert actual == expected, f"{path}: {actual!r} != {expected!r}" + + +def assert_storage_data_exact(actual, expected, *, context): + """Assert every tensor-valued data, scale, and amax metadata field exactly.""" + assert_nested_state_exact( + snapshot_storage_tensor_metadata(actual), + snapshot_storage_tensor_metadata(expected), + path=context, + ) + + +def assert_hybrid_tensor_exact(actual, expected, *, context): + """Compare both hybrid directions, including metadata and dequantization.""" + for direction in ("rowwise", "columnwise"): + actual_storage = getattr(actual, f"{direction}_sub_storage") + expected_storage = getattr(expected, f"{direction}_sub_storage") + assert_storage_data_exact( + actual_storage, + expected_storage, + context=f"{context} {direction}", + ) + if expected_storage is None: + continue + try: + expected_dequantized = expected_storage.dequantize() + except NotImplementedError: + continue + torch.testing.assert_close( + actual_storage.dequantize(), + expected_dequantized, + rtol=0.0, + atol=0.0, + msg=f"{context} {direction} dequantized value differs", + ) + + +def make_role_aware_quantizer(factory, role): + """Construct a quantizer with the standard Block-FP8 GEMM geometry.""" + quantizer = factory() + if isinstance(quantizer, te.Float8BlockQuantizer): + is_weight = ( + role is not None + and role.module_type in _LINEAR_MODULE_TYPES + and role.tensor_type == "weight" + ) + quantizer.block_scaling_dim = 2 if is_weight else 1 + return quantizer + + +def hybrid_custom_recipe(row_factory, col_factory, grad_factory=None): + """Build a CustomRecipe with hybrid forward and configurable grad quantizers.""" + if grad_factory is None: + grad_factory = col_factory + + def qfactory(role): + is_linear = _is_linear_role(role) + if is_linear and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=make_role_aware_quantizer(row_factory, role), + columnwise_quantizer=make_role_aware_quantizer(col_factory, role), + ) + if is_linear and role.tensor_type in _GRAD_TENSOR_TYPES: + return make_role_aware_quantizer(grad_factory, role) + return make_role_aware_quantizer(row_factory, role) + + return recipe.CustomRecipe(qfactory=qfactory) + + +def hybrid_fp8_current_qfactory(role): + """FP8 current scaling in both hybrid directions for forward tensor roles.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=current_scaling_factory(role), + columnwise_quantizer=current_scaling_factory(role), + ) + return current_scaling_factory(role) + + +def hybrid_fp8_current_e5m2_grads_qfactory(role): + """FP8 current-scaling hybrid with E5M2 for both grad boundary roles.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=_make_fp8_current(), + columnwise_quantizer=_make_fp8_current(), + ) + if _is_linear_role(role) and role.tensor_type in _GRAD_TENSOR_TYPES: + return _make_fp8_current(fp8_dtype=te.DType.kFloat8E5M2) + return _make_fp8_current() + + +def hybrid_mxfp8_qfactory(role): + """MXFP8 in both hybrid directions for forward tensor roles.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=mxfp8_factory(role), + ) + return mxfp8_factory(role) + + +def hybrid_float8_block_qfactory(role): + """Float8 block scaling in both hybrid directions for forward roles.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=float8_block_scaling_factory(role), + columnwise_quantizer=float8_block_scaling_factory(role), + ) + return float8_block_scaling_factory(role) + + +def hybrid_block_fp8_e4m3_qfactory(role): + """Hybrid E4M3 Block-FP8 with role-aware 1D/2D block geometry.""" + is_linear = _is_linear_role(role) + is_weight = is_linear and role.tensor_type == "weight" + block_scaling_dim = 2 if is_weight else 1 + + def make_quantizer(): + return te.Float8BlockQuantizer( + fp8_dtype=te.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=block_scaling_dim, + ) + + if is_linear and role.tensor_type in _GRAD_TENSOR_TYPES: + return make_quantizer() + return te.HybridQuantizer( + rowwise_quantizer=make_quantizer(), + columnwise_quantizer=make_quantizer(), + ) + + +def hybrid_mixed_mxfp8_fp8_qfactory(role): + """MXFP8 rowwise plus FP8 current-scaling columnwise.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=current_scaling_factory(role), + ) + return current_scaling_factory(role) + + +def hybrid_fp8_current_identity_qfactory(role): + """FP8 current-scaling forward plus Identity backward.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=current_scaling_factory(role), + columnwise_quantizer=te.IdentityQuantizer(), + ) + if _is_linear_role(role) and role.tensor_type in _GRAD_TENSOR_TYPES: + return te.IdentityQuantizer() + return current_scaling_factory(role) + + +def hybrid_mxfp8_identity_qfactory(role): + """MXFP8 forward plus Identity backward.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=te.IdentityQuantizer(), + ) + if _is_linear_role(role) and role.tensor_type in _GRAD_TENSOR_TYPES: + return te.IdentityQuantizer() + return mxfp8_factory(role) + + +def identity_qfactory(role): # pylint: disable=unused-argument + """High-precision passthrough for every quantizer slot.""" + return te.IdentityQuantizer() + + +def hybrid_nvfp4_qfactory(role): + """NVFP4 in both hybrid directions for forward tensor roles.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=nvfp4_factory(role), + columnwise_quantizer=nvfp4_factory(role), + ) + return nvfp4_factory(role) + + +def hybrid_tp_mxfp8_nvfp4_qfactory(role): + """TP/SP MXFP8 rowwise plus NVFP4 columnwise, including boundary roles.""" + if _is_linear_role(role) and role.tensor_type in _GRAD_TENSOR_TYPES: + return nvfp4_factory(role) + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=nvfp4_factory(role), + ) + + +def hybrid_fp8_mxfp8_qfactory(role): + """FP8 current-scaling rowwise plus MXFP8 columnwise for CPU-offload tests.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=_make_fp8_current(), + columnwise_quantizer=_make_mxfp8(), + ) + if _is_linear_role(role) and role.tensor_type in _GRAD_TENSOR_TYPES: + return _make_mxfp8(fp8_dtype=te.DType.kFloat8E5M2) + return _make_fp8_current() + + +def hybrid_mxfp8_nvfp4_qfactory(role): + """MXFP8 rowwise plus NVFP4 columnwise for CPU-offload tests.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=_make_mxfp8(), + columnwise_quantizer=nvfp4_factory(role), + ) + if _is_linear_role(role) and role.tensor_type in _GRAD_TENSOR_TYPES: + return nvfp4_factory(role) + return _make_mxfp8() diff --git a/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py b/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py index 306d0627f5..34aaf32ec9 100644 --- a/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py +++ b/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py @@ -137,10 +137,41 @@ def _param_key(name): return name.split(".")[-1] +def _no_checkpoint_activation_bytes(cfg, seq_size, itemsize): + """Activations LayerNormMLP saves for backward when checkpoint=False. + + Per layer: ln_out and out (seq*hidden each), fc1_out and act_out + (seq*ffn_hidden each), mu and rsigma (seq each). + """ + per_layer = 2 * seq_size * (cfg._ffn_hidden_size + cfg._hidden_size + 1) + return cfg._layers * per_layer * itemsize + + +def _recomputed_activation_bytes(cfg, seq_size, itemsize): + """Activations checkpointing must free: fc1_out and act_out. + + The peak still holds the transient of one layer, so only the remaining + layers count. Keeping this independent of _layers means the assertion + below does not encode the shape of the test models. + """ + return (cfg._layers - 1) * 2 * seq_size * cfg._ffn_hidden_size * itemsize + + @pytest.mark.parametrize("size", config.keys()) @pytest.mark.parametrize("seq_size", seq_sizes) def test_selective_activation_checkpoint(size, seq_size): + itemsize = torch.empty((), dtype=torch.get_default_dtype()).element_size() + no_ckpt_bytes = _no_checkpoint_activation_bytes(config[size], seq_size, itemsize) + + # Both models live in the same process, so budget the non-checkpointed peak twice. + free_bytes, _ = torch.cuda.mem_get_info(device) + if free_bytes < 2 * no_ckpt_bytes: + pytest.skip( + f"needs {2 * no_ckpt_bytes / 2**30:.1f} GiB free device memory, only" + f" {free_bytes / 2**30:.1f} GiB available" + ) + ln_model, sln_model = config[size].build() data = torch.randn((seq_size, config[size]._hidden_size), device=device) @@ -152,15 +183,8 @@ def test_selective_activation_checkpoint(size, seq_size): sln_fwd_out, sln_fwd_time, sln_fwd_mem = _run_fwd(sln_model, data) sln_grads, sln_bwd_time, sln_bwd_mem = _run_bwd(sln_model, sln_fwd_out) - assert ln_fwd_mem > 6 * sln_fwd_mem, ( - "selective activation checkpointing does not reduce forward memory by 6X, only by" - f" {ln_fwd_mem/sln_fwd_mem}!" - ) - assert ln_bwd_time < sln_bwd_time, ( - "selective activation activation checkpointing backward pass is NOT slower than native!" - f" got Native LayerNormMLP Backward Time: {ln_bwd_time} ms and Selective Activation" - f" Checkpointed LayerNormMLP Backward Time: {sln_bwd_time} ms" - ) + # Correctness first, so that a numerical regression is not masked by the + # memory check below. diff = _max_diff(ln_fwd_out, sln_fwd_out) assert diff == 0.0, f"outputs are not equal! maximum difference {diff}" for key in [ @@ -173,3 +197,12 @@ def test_selective_activation_checkpoint(size, seq_size): ]: diff = _max_diff(ln_grads[key], sln_grads[key]) assert diff == 0.0, f"gradients for {key} are not equal! maximum difference: {diff}" + + # Checkpointing recomputes fc1_out and act_out, so it must free at least those. + expected_saving = _recomputed_activation_bytes(config[size], seq_size, itemsize) + saving = ln_fwd_mem - sln_fwd_mem + assert saving >= 0.95 * expected_saving, ( + "selective activation checkpointing did not free the recomputed activations: saved" + f" {saving} B, expected at least {0.95 * expected_saving} B (ln_fwd_mem={ln_fwd_mem}," + f" sln_fwd_mem={sln_fwd_mem})" + ) diff --git a/tests/pytorch/mxfp8/test_mxfp8_dequantize_extreme_scales.py b/tests/pytorch/mxfp8/test_mxfp8_dequantize_extreme_scales.py new file mode 100644 index 0000000000..64bfe6012e --- /dev/null +++ b/tests/pytorch/mxfp8/test_mxfp8_dequantize_extreme_scales.py @@ -0,0 +1,28 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import pytest +import torch + +import transformer_engine.pytorch as te +from transformer_engine.pytorch import MXFP8Quantizer + +recipe_available, reason_for_no_recipe = te.is_mxfp8_available(return_reason=True) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_dequantize_extreme_e8m0_scale_codes() -> None: + """UE8M0 scale code 0 is 2^-127 and code 255 is NaN, not 0.0 and +Inf.""" + quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, columnwise=False) + x = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") + qx = quantizer(x) + data = qx._rowwise_data.view(torch.uint8) + scales = qx._rowwise_scale_inv.view(torch.uint8) + data[0, :32] = 56 + scales[0, 0] = 0 + data[1, :32] = 56 + scales[1, 0] = 255 + y = qx.dequantize(dtype=torch.float32) + assert (y[0, :32] == 2.0**-127).all() + assert torch.isnan(y[1, :32]).all() diff --git a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py index d07953ce37..c0cb0aa09a 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py +++ b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py @@ -60,6 +60,13 @@ def generate_split_sections(M: int, N: int, edge_cases: str) -> list[int]: split_sections = [avg_split] * (num_chunks - 2) + [0] + [avg_split * 2] elif edge_cases == "random_uneven_split": split_sections = generate_random_multiples_sum(M, num_chunks, least_multiple) + elif edge_cases == "imbalanced_avg_misaligned": + # Three groups whose 128-aligned sizes average to a non-128-aligned value, so any + # buffer sized from num_groups * round_up(M // num_groups, 128) diverges from the + # per-group padded sum. "random_uneven_split" cannot catch that: with 4 chunks the + # average only depends on M, which is always 128 * num_chunks aligned here. + assert M >= 3 * least_multiple, "M too small for the imbalanced case" + split_sections = [least_multiple, least_multiple, M - 2 * least_multiple] else: raise ValueError(f"Invalid edge case: {edge_cases}") @@ -137,6 +144,7 @@ def check_grouped_tensor_mxfp8_versus_reference( return_transpose: bool, split_sections: list[int], optimize_for_gemm: bool = False, + with_2d_quantization: bool = False, ) -> None: te_dtype = te.DType.kFloat8E4M3 @@ -159,6 +167,7 @@ def check_grouped_tensor_mxfp8_versus_reference( fp8_dtype=te_dtype, rowwise=return_rowwise, columnwise=return_transpose, + with_2d_quantization=with_2d_quantization, ) for _ in range(len(split_sections)) ] @@ -330,6 +339,30 @@ def check_grouped_tensor_mxfp8_with_paged_stashing( torch.testing.assert_close(x_sx_t_i, x_sx_t_ref_i, atol=0.0, rtol=0.0) +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("quantize_mode", ["rowwise_only", "both_directions", "columnwise_only"]) +@pytest.mark.parametrize( + "optimize_for_gemm", [True, False], ids=["optimize_for_gemm", "no_optimize_for_gemm"] +) +def test_grouped_tensor_mxfp8_2d_quantization_versus_reference( + quantize_mode: str, + optimize_for_gemm: bool, +) -> None: + """Grouped MXFP8 should match independent 2D quantization of each tensor.""" + return_rowwise = quantize_mode != "columnwise_only" + return_transpose = quantize_mode != "rowwise_only" + check_grouped_tensor_mxfp8_versus_reference( + x_dtype=torch.bfloat16, + M=1024, + N=256, + return_rowwise=return_rowwise, + return_transpose=return_transpose, + split_sections=[256, 256, 256, 256], + optimize_for_gemm=optimize_for_gemm, + with_2d_quantization=True, + ) + + @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.parametrize( "M, N", @@ -469,3 +502,419 @@ def test_grouped_tensor_mxfp8_with_paged_stashing( valid_M=valid_M, optimize_for_gemm=optimize_for_gemm, ) + + +# --------------------------------------------------------------------------------------------- +# Pre-quantized MXFP8 input (FP8 token dispatch) +# +# tex.group_requantize_inplace takes a grouped tensor that arrives +# ALREADY rowwise-quantized (its high-precision form no longer exists), and makes it GEMM-ready +# in both directions: the rowwise data passes through verbatim with its scales swizzled, and the +# columnwise copy is rebuilt via dequantize + columnwise-only requantize. +# +# These mirror the edge-case matrices of test_grouped_tensor_mxfp8_versus_reference and +# test_grouped_tensor_mxfp8_with_paged_stashing so the same shapes, zero-token placements and +# uneven splits exercise this path. +# --------------------------------------------------------------------------------------------- + + +def make_prequantized_wire_tensor(x: torch.Tensor, split_section_tensor: torch.Tensor): + """Rowwise-only, unswizzled grouped tensor, as FP8 dispatch delivers it.""" + wire_quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, rowwise=True, columnwise=False) + # Must stay unswizzled: the requantize path asserts on it and dequantize needs compact scales. + wire_quantizer.optimize_for_gemm = False + wire = fused_grouped_quantize(x, split_section_tensor, wire_quantizer) + assert wire.columnwise_data is None + assert not wire._with_gemm_swizzled_scales + return wire + + +def make_op_quantizer(columnwise: bool): + """The op's input quantizer, configured the way the ops layer configures it. + + ``columnwise`` mirrors ``weight_requires_grad``: it tells the helper whether a wgrad GEMM + will consume a columnwise copy. + """ + quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, rowwise=True, columnwise=columnwise) + quantizer.optimize_for_gemm = True + return quantizer + + +def make_gemm_ready_tensor(x: torch.Tensor, split_section_tensor: torch.Tensor): + """Grouped tensor already GEMM-ready in both directions (swizzled scales).""" + quantizer = make_op_quantizer(columnwise=True) + tensor = fused_grouped_quantize(x, split_section_tensor, quantizer) + assert tensor.columnwise_data is not None + assert tensor._with_gemm_swizzled_scales + return tensor + + +def check_prequantized_requantize_versus_reference( + x_dtype: torch.dtype, + M: int, + N: int, + split_sections: list[int], +) -> None: + """Run the pre-quantized requantize path and check both directions against a reference. + + The reference is derived from dequantize(wire), not from the original high-precision x: that + is the only data a consumer can see after dispatch, and MXFP8 rowwise requantization is + idempotent, so it is an exact reference rather than an approximate one. + """ + device = "cuda" + torch.manual_seed(0) + torch.cuda.manual_seed(0) + + # The buffer is always M rows. Paged stashing is just the case where the groups cover fewer + # than M of them (valid_M < M) and the tail holds garbage the kernels must leave alone; the + # non-paged case is the same code path with sum(split_sections) == M. + x = torch.randn((M, N), dtype=x_dtype, device=device) + split_section_tensor = torch.tensor(split_sections, dtype=torch.int64, device=device) + num_groups = len(split_sections) + # Rows the groups actually cover. Beyond this the buffers hold whatever the allocator handed + # out, so nothing past it may be compared. + valid_rows = sum(split_sections) + + wire = make_prequantized_wire_tensor(x, split_section_tensor) + + # Snapshot what must survive verbatim, plus the compact scales the reference swizzles. + rowwise_scale_shape_before = wire.scale_inv.shape + rowwise_data_before = wire.rowwise_data.clone() + wire_splits_before = [ + (t._rowwise_data.view(dtype=torch.uint8).clone(), t._rowwise_scale_inv.clone()) + for t in wire.split_into_quantized_tensors() + ] + + # Reference high-precision input: everything downstream is derived from this. Only the live + # rows are kept -- dequantize allocates M rows but writes only the ones the groups cover. + dequantized_ref = ( + tex.group_dequantize(wire, te.DType.kBFloat16) + .rowwise_data.view(M, N)[:valid_rows, :] + .clone() + ) + + # Reference columnwise copy, quantized per group from the dequantized data. + colwise_quantizers = [ + MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, rowwise=False, columnwise=True) + for _ in range(num_groups) + ] + _, _, colwise_data_ref, colwise_scale_ref = reference_group_quantize( + dequantized_ref, + colwise_quantizers, + split_sections, + return_rowwise=False, + return_transpose=True, + ) + + # ---- the code under test ---- + # The op's quantizer, configured as the ops layer does: columnwise_usage says a wgrad GEMM + # will consume the columnwise copy, so the helper builds it and switches rowwise off itself. + dequantized = tex.group_requantize_inplace( + wire, + make_op_quantizer(columnwise=True), + num_groups, + split_section_tensor, + te.DType.kBFloat16, + return_dequantized=True, + ) + + assert wire.columnwise_data is not None, "columnwise data must be built" + assert wire.columnwise_scale_inv is not None, "columnwise scales must survive the swizzle" + + # The rowwise swizzle is size-preserving; the per-group content checks below cannot catch a + # wrongly sized buffer because they read through offsets derived from the splits, so the + # capacity must be checked explicitly. + assert ( + wire.scale_inv.shape == rowwise_scale_shape_before + ), "the swizzled rowwise scale buffer must keep the compact buffer's capacity" + + # The returned dequantized tensor is what bias gradients are reduced from. Compare only the + # live rows: both this and the reference allocate M rows but write only the covered ones, and + # their tails are separate uninitialized allocations. + torch.testing.assert_close(dequantized[:valid_rows, :], dequantized_ref, atol=0.0, rtol=0.0) + + # The rowwise DATA must pass through untouched; only its scales are re-laid-out. + torch.testing.assert_close(wire.rowwise_data, rowwise_data_before, atol=0.0, rtol=0.0) + + if valid_rows > 0: + # A tensor whose groups are all empty has no scales to lay out, so the swizzle is a no-op + # and leaves the flag unset; every other case must come back swizzled. + assert wire._with_gemm_swizzled_scales, "rowwise scales must be marked swizzled" + + # Per-group comparison, same structure as check_grouped_tensor_mxfp8_versus_reference. + outputs = wire.split_into_quantized_tensors() + x_splits = torch.split(dequantized_ref, split_sections) + + for i, out in enumerate(outputs): + rows_i = split_sections[i] + scale_before = wire_splits_before[i][1] + colwise_data = out._columnwise_data.view(dtype=torch.uint8) + colwise_scale = out._columnwise_scale_inv + + if rows_i == 0: + # Buffers for empty groups are never written, so only shape and dtype are meaningful. + assert_same_shape_and_dtype(colwise_data, colwise_data_ref[i]) + assert_same_shape_and_dtype(colwise_scale, colwise_scale_ref[i]) + continue + + # Rowwise scales: the swizzled form of the compact scales this group arrived with. The + # rowwise DATA is covered by the whole-buffer identity check above. + torch.testing.assert_close( + out._rowwise_scale_inv, + swizzle_mxfp8_scale(rows_i, N, scale_before, columnwise=False), + atol=0.0, + rtol=0.0, + ) + + # Columnwise: rebuilt from the dequantized data, and swizzled by the quantize kernel + # because the caller sets optimize_for_gemm. + torch.testing.assert_close(colwise_data, colwise_data_ref[i], atol=0.0, rtol=0.0) + valid_scale_shape = get_mxfp8_scale_shape_no_padding(x_splits[i].shape, True) + assert ( + valid_scale_shape == colwise_scale.shape + ), "The columnwise scale shape is not correctly aligned" + torch.testing.assert_close( + colwise_scale, + swizzle_mxfp8_scale(rows_i, N, colwise_scale_ref[i], columnwise=True), + atol=0.0, + rtol=0.0, + ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "M, N", + [ + # edge case, zero tokens for all + (0, 512), + # full tile cases + (1024, 256), + # larger sizes + (8192, 1024), + (16384, 8192), + ], +) +@pytest.mark.parametrize("x_dtype", [torch.bfloat16], ids=str) +@pytest.mark.parametrize( + "edge_cases", + [ + "regular", + "zero_tokens_front", + "zero_tokens_end", + "zero_tokens_middle", + "random_uneven_split", + "imbalanced_avg_misaligned", + ], +) +def test_prequantized_requantize_versus_reference( + x_dtype: torch.dtype, + M: int, + N: int, + edge_cases: str, +) -> None: + split_sections = generate_split_sections(M, N, edge_cases) + check_prequantized_requantize_versus_reference( + x_dtype=x_dtype, + M=M, + N=N, + split_sections=split_sections, + ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "M, N", + [ + # M won't be empty in paged stashing + (1024, 256), + (8192, 1024), + (16384, 8192), + ], +) +@pytest.mark.parametrize("x_dtype", [torch.bfloat16], ids=str) +@pytest.mark.parametrize( + "edge_cases", + [ + "regular", + "zero_tokens_all", + "zero_tokens_front", + "zero_tokens_end", + "zero_tokens_middle", + "random_uneven_split", + "imbalanced_avg_misaligned", + ], +) +def test_prequantized_requantize_with_paged_stashing( + x_dtype: torch.dtype, + M: int, + N: int, + edge_cases: str, +) -> None: + # Paged stashing: the buffer holds M rows but only valid_M carry live tokens; the rest is + # garbage the kernels must not touch. + valid_M = 0 if edge_cases == "zero_tokens_all" else M // 2 + split_sections = generate_split_sections(valid_M, N, edge_cases) + assert sum(split_sections) == valid_M + + check_prequantized_requantize_versus_reference( + x_dtype=x_dtype, + M=M, + N=N, + split_sections=split_sections, + ) + + +def _requantize_setup(M: int = 1024, N: int = 256): + """Common inputs for the state-dispatch tests below.""" + torch.manual_seed(0) + split_sections = [M // 4] * 4 + x = torch.randn((M, N), dtype=torch.bfloat16, device="cuda") + return x, torch.tensor(split_sections, dtype=torch.int64, device="cuda"), len(split_sections) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_prequantized_requantize_passes_through_gemm_ready_input(): + """A tensor already GEMM-ready in both directions is left untouched.""" + x, splits, num_groups = _requantize_setup() + tensor = make_gemm_ready_tensor(x, splits) + rowwise_before = tensor.rowwise_data.clone() + columnwise_before = tensor.columnwise_data.clone() + scale_before = tensor.scale_inv.clone() + + out = tex.group_requantize_inplace( + tensor, make_op_quantizer(columnwise=True), num_groups, splits, te.DType.kBFloat16 + ) + + assert out is None + assert tensor._with_gemm_swizzled_scales + torch.testing.assert_close(tensor.rowwise_data, rowwise_before, atol=0.0, rtol=0.0) + torch.testing.assert_close(tensor.columnwise_data, columnwise_before, atol=0.0, rtol=0.0) + torch.testing.assert_close(tensor.scale_inv, scale_before, atol=0.0, rtol=0.0) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_prequantized_requantize_skips_columnwise_when_not_needed(): + """columnwise_usage=False (frozen weights) swizzles rowwise without building columnwise.""" + x, splits, num_groups = _requantize_setup() + wire = make_prequantized_wire_tensor(x, splits) + + out = tex.group_requantize_inplace( + wire, make_op_quantizer(columnwise=False), num_groups, splits, te.DType.kBFloat16 + ) + + assert out is None + assert wire._with_gemm_swizzled_scales, "the GEMM still needs swizzled rowwise scales" + assert wire.columnwise_data is None, "no wgrad GEMM, so no columnwise copy should be built" + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_prequantized_requantize_rejects_dequantized_from_gemm_ready_input(): + """Bias grads cannot be served from an already-swizzled input, so this must raise.""" + x, splits, num_groups = _requantize_setup() + tensor = make_gemm_ready_tensor(x, splits) + + with pytest.raises(RuntimeError, match="compact format"): + tex.group_requantize_inplace( + tensor, + make_op_quantizer(columnwise=True), + num_groups, + splits, + te.DType.kBFloat16, + return_dequantized=True, + ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_prequantized_requantize_rejects_swizzled_without_columnwise(): + """Swizzled rowwise scales with no columnwise copy: it can no longer be rebuilt.""" + x, splits, num_groups = _requantize_setup() + wire = make_prequantized_wire_tensor(x, splits) + # Swizzle in place, leaving the tensor rowwise-only. + tex.grouped_swizzle_for_gemm(wire, True, False) + + with pytest.raises(RuntimeError, match="cannot be rebuilt"): + tex.group_requantize_inplace( + wire, make_op_quantizer(columnwise=True), num_groups, splits, te.DType.kBFloat16 + ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_prequantized_requantize_rejects_dtype_mismatch(): + """The helper keeps the input's format; it does not convert between formats.""" + x, splits, num_groups = _requantize_setup() + wire = make_prequantized_wire_tensor(x, splits) + mismatched = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E5M2, rowwise=True, columnwise=True) + mismatched.optimize_for_gemm = True + + with pytest.raises(RuntimeError, match="dtype"): + tex.group_requantize_inplace(wire, mismatched, num_groups, splits, te.DType.kBFloat16) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("columnwise", [False, True], ids=["rowwise", "columnwise"]) +@pytest.mark.parametrize( + "split_sections", + [ + # Imbalanced groups whose mean is not 128-aligned: sizing the swizzled output from + # num_tensors * round_up(mean, 128) would give 3 * 512 = 1536 rows instead of the + # input's 1280. + [128, 128, 1024], + # Same property with a zero-token group in the mix. + [0, 256, 128, 1152], + ], + ids=["imbalanced", "imbalanced_with_empty"], +) +def test_grouped_swizzle_variable_shape_preserves_scale_capacity( + columnwise: bool, split_sections: list[int] +): + """Swizzling a variable-shape grouped tensor must keep the scale buffer's exact capacity. + + The swizzle kernel walks input and output with identical per-group padded strides, so the + operation is size-preserving. Sizing the output from the per-tensor average instead breaks + any consumer that derives per-group offsets from the split sizes, because the average of + 128-aligned group sizes is generally not 128-aligned. The content checks alone cannot catch + a wrong allocation (per-group offsets are derived from the splits, not the buffer), so the + shape assertion is the actual regression check. + """ + torch.manual_seed(0) + N = 256 + M = sum(split_sections) + x = torch.randn((M, N), dtype=torch.bfloat16, device="cuda") + splits = torch.tensor(split_sections, dtype=torch.int64, device="cuda") + + quantizer = MXFP8Quantizer( + fp8_dtype=te.DType.kFloat8E4M3, rowwise=not columnwise, columnwise=columnwise + ) + # Compact (unswizzled) scales, as they arrive over the wire. + quantizer.optimize_for_gemm = False + tensor = fused_grouped_quantize(x, splits, quantizer) + assert not tensor._with_gemm_swizzled_scales + + scale_attr = "columnwise_scale_inv" if columnwise else "scale_inv" + compact_shape = getattr(tensor, scale_attr).shape + compact_groups = [ + (t._columnwise_scale_inv if columnwise else t._rowwise_scale_inv).clone() + for t in tensor.split_into_quantized_tensors() + ] + + tex.grouped_swizzle_for_gemm(tensor, not columnwise, columnwise) + + assert tensor._with_gemm_swizzled_scales + assert ( + getattr(tensor, scale_attr).shape == compact_shape + ), "grouped swizzle must preserve the scale buffer's shape for variable-shape tensors" + + # The swizzled content of each group must match a per-group dense swizzle of the compact + # scales it arrived with. + for rows_i, group_before, out in zip( + split_sections, compact_groups, tensor.split_into_quantized_tensors() + ): + if rows_i == 0: + continue + out_scale = out._columnwise_scale_inv if columnwise else out._rowwise_scale_inv + torch.testing.assert_close( + out_scale, + swizzle_mxfp8_scale(rows_i, N, group_before, columnwise=columnwise), + atol=0.0, + rtol=0.0, + ) diff --git a/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py b/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py index 127b487650..16a2d75de6 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py +++ b/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py @@ -129,3 +129,36 @@ def test_mxfp8_quantize_swizzle_fusion( return_rowwise=return_rowwise, return_transpose=return_transpose, ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("M, N", [(96, 160), (4096, 576), (4096, 2112)]) +def test_mxfp8_bidirectional_swizzled_row_scale_padding(M: int, N: int) -> None: + """The specialized bidirectional kernel must not overwrite padded row scales.""" + x = torch.randn((M, N), dtype=torch.bfloat16, device="cuda") + quantizer = MXFP8Quantizer( + fp8_dtype=te.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + ) + quantizer.optimize_for_gemm = True + scale = quantizer(x)._rowwise_scale_inv.view(torch.uint8) + + scale_rows = torch.arange(M, device=scale.device, dtype=torch.int64).view(-1, 1) + scale_cols = torch.arange(N // 32, device=scale.device, dtype=torch.int64).view(1, -1) + num_tiles_x = math.ceil(N / 128) + scale_indices = ( + ((scale_rows // 128) * num_tiles_x + scale_cols // 4) * (128 * 4) + + (scale_rows % 32) * 16 + + ((scale_rows % 128) // 32) * 4 + + scale_cols % 4 + ) + valid_mask = torch.zeros(scale.numel(), dtype=torch.bool, device=scale.device) + valid_mask[scale_indices.view(-1)] = True + + torch.testing.assert_close( + scale.view(-1)[~valid_mask], + torch.zeros_like(scale.view(-1)[~valid_mask]), + atol=0, + rtol=0, + ) diff --git a/tests/pytorch/test_nvfp4_fsdp2_hooks.py b/tests/pytorch/nvfp4/test_nvfp4_fsdp2_hooks.py similarity index 100% rename from tests/pytorch/test_nvfp4_fsdp2_hooks.py rename to tests/pytorch/nvfp4/test_nvfp4_fsdp2_hooks.py diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index b3ef196bf3..60139700ef 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -11,8 +11,9 @@ from transformer_engine.pytorch.constants import TE_DType from transformer_engine.pytorch import NVFP4Quantizer from transformer_engine.pytorch.cpp_extensions import general_gemm, general_grouped_gemm -from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef -from transformer_engine.pytorch.custom_recipes import utils +from transformer_engine.pytorch.custom_recipes.reference_nvfp4 import NVFP4QuantizerRef +from transformer_engine.pytorch.custom_recipes import reference_utils +from transformer_engine.pytorch.tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage from torch.utils.cpp_extension import IS_HIP_EXTENSION if IS_HIP_EXTENSION: @@ -137,7 +138,7 @@ def check_nvfp4_gemm_versus_reference( # Create reference quantizer for reference GEMM x_ref_quantizer = NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + dtype=reference_utils.Fp4Formats.E2M1, rowwise=True, columnwise=not row_scaled_nvfp4, pow_2_scales=False, @@ -149,7 +150,7 @@ def check_nvfp4_gemm_versus_reference( nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) w_ref_quantizer = NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + dtype=reference_utils.Fp4Formats.E2M1, rowwise=True, columnwise=True, pow_2_scales=False, @@ -434,6 +435,65 @@ def check_nvfp4_row_scaled_gemm_matches_emulated( torch.testing.assert_close(y_row_scaled, y_emulated, atol=3.0517578125e-5, rtol=0.0) +def _dequantize_nvfp4_usage( + tensor: NVFP4TensorStorage, + *, + columnwise: bool, +) -> torch.Tensor: + """Dequantize one independently-quantized tensor orientation.""" + if not columnwise: + return tensor.dequantize(dtype=torch.float32) + + metadata = tensor.get_metadata() + metadata["rowwise_data"] = metadata["columnwise_data"] + metadata["rowwise_scale_inv"] = metadata["columnwise_scale_inv"] + metadata["amax_rowwise"] = metadata["amax_columnwise"] + metadata["columnwise_data"] = None + metadata["columnwise_scale_inv"] = None + metadata["amax_columnwise"] = None + return NVFP4TensorStorage(**metadata).dequantize(dtype=torch.float32) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "layout,a_shape,b_shape", + [ + ("TN", (64, 64), (96, 64)), + ("NN", (64, 96), (128, 64)), + ("NT", (128, 64), (128, 96)), + ], +) +def test_nvfp4_bilateral_row_scaled_gemm_matches_dequantized( + layout: str, + a_shape: tuple[int, int], + b_shape: tuple[int, int], +) -> None: + """Check per-tensor GEMM plus bilateral FP32 post-scaling.""" + torch.manual_seed(41) + quantizer = NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=False, + with_post_rht_amax=False, + row_scaled_nvfp4=True, + ) + a = torch.randn(a_shape, dtype=torch.bfloat16, device="cuda") + b = torch.randn(b_shape, dtype=torch.bfloat16, device="cuda") + a_nvfp4 = quantizer(a) + b_nvfp4 = quantizer(b) + + actual = general_gemm(a_nvfp4, b_nvfp4, out_dtype=torch.float32, layout=layout)[0] + transa, transb = layout[0] == "T", layout[1] == "T" + a_dequant = _dequantize_nvfp4_usage(a_nvfp4, columnwise=not transa) + b_dequant = _dequantize_nvfp4_usage(b_nvfp4, columnwise=transb) + expected = b_dequant @ a_dequant.T + + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.parametrize( "M, K, N", diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py index a068db7623..98776f6bd7 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py @@ -15,8 +15,8 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex from transformer_engine.pytorch import NVFP4Quantizer -from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef -from transformer_engine.pytorch.custom_recipes import utils +from transformer_engine.pytorch.custom_recipes.reference_nvfp4 import NVFP4QuantizerRef +from transformer_engine.pytorch.custom_recipes import reference_utils from transformer_engine.common.recipe import NVFP4BlockScaling import pytest @@ -259,3 +259,25 @@ def test_rht_with_quantization_block_tiling_versus_reference( with_random_sign_mask=with_random_sign_mask, optimize_for_gemm=optimize_for_gemm, ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("quantize_mode", ["rowwise_only", "both_directions", "columnwise_only"]) +def test_rht_split_quantize_matches_per_tensor_reference(quantize_mode: str) -> None: + # split_quantize sends RHT-enabled NVFP4 quantizers to the grouped Hadamard + # transform kernels, which are implemented for the SM100 family only. On + # other architectures it falls back to quantizing each split on its own. + # Both routes have to give the same result as per-tensor quantization. + split_sections = [128, 128, 128, 128] + return_rowwise = quantize_mode in ("rowwise_only", "both_directions") + return_transpose = quantize_mode in ("columnwise_only", "both_directions") + + check_group_quantization_nvfp4_versus_reference( + x_dtype=torch.bfloat16, + M=sum(split_sections), + N=256, + return_rowwise=return_rowwise, + return_transpose=return_transpose, + split_sections=split_sections, + with_rht=True, + ) diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py b/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py index 8a5a95fdc1..aefd371e4e 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py @@ -8,8 +8,8 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex from transformer_engine.pytorch import NVFP4Quantizer -from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef -from transformer_engine.pytorch.custom_recipes import utils +from transformer_engine.pytorch.custom_recipes.reference_nvfp4 import NVFP4QuantizerRef +from transformer_engine.pytorch.custom_recipes import reference_utils from transformer_engine.common.recipe import NVFP4BlockScaling from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor from torch.utils.cpp_extension import IS_HIP_EXTENSION diff --git a/tests/pytorch/nvfp4/test_nvfp4_module_exact.py b/tests/pytorch/nvfp4/test_nvfp4_module_exact.py index b57b78eb13..9fdabf68dd 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_module_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_module_exact.py @@ -6,8 +6,8 @@ import torch import transformer_engine.pytorch as te from transformer_engine.common import recipe -from transformer_engine.pytorch.custom_recipes import quantization_ref_nvfp4 -from transformer_engine.pytorch.custom_recipes import utils +from transformer_engine.pytorch.custom_recipes import reference_nvfp4 +from transformer_engine.pytorch.custom_recipes import reference_utils recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) @@ -53,7 +53,23 @@ def nvfp4_rht_and_2d_quantization(): return nvfp4_recipe @staticmethod - def nvfp4_recipe_to_test(with_rht: bool = False, with_2d_quantization: bool = False): + def nvfp4_row_scaled(): + nvfp4_recipe = recipe.NVFP4BlockScaling() + nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() + nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams() + nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams() + # Emit row-scaled (per-token) NVFP4 for the forward activation. In a + # training Linear this also drives the row-scaled columnwise (transpose) + # of the activation that the wgrad GEMM consumes in the backward pass. + nvfp4_recipe.row_scaled_activation = True + return nvfp4_recipe + + @staticmethod + def nvfp4_recipe_to_test( + with_rht: bool = False, with_2d_quantization: bool = False, row_scaled: bool = False + ): + if row_scaled: + return GetRecipes.nvfp4_row_scaled() if with_rht and with_2d_quantization: return GetRecipes.nvfp4_rht_and_2d_quantization() elif with_rht: @@ -64,7 +80,9 @@ def nvfp4_recipe_to_test(with_rht: bool = False, with_2d_quantization: bool = Fa return GetRecipes.nvfp4_vanilla() -def get_nvfp4_quantizer_factory(with_rht: bool = False, with_2d_quantization: bool = False): +def get_nvfp4_quantizer_factory( + with_rht: bool = False, with_2d_quantization: bool = False, row_scaled: bool = False +): """ Create a quantizer factory for NVFP4 reference implementation. @@ -85,8 +103,8 @@ def get_nvfp4_quantizer_factory(with_rht: bool = False, with_2d_quantization: bo # qfactory, so we return a valid quantizer for those slots; it is harmless because # the GEMM outputs in the high-precision activation dtype, not in NVFP4. def _default_quantizer(): - return quantization_ref_nvfp4.NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + return reference_nvfp4.NVFP4QuantizerRef( + dtype=reference_utils.Fp4Formats.E2M1, quant_tile_shape=(1, 16), pow_2_scales=False, with_rht=with_rht, @@ -96,10 +114,18 @@ def factory(role): if role is None: return _default_quantizer() if role.tensor_type == "input": - return _default_quantizer() + # Only the forward activation is row-scaled, mirroring the production + # wiring in quantization.py (mode == "forward" and tensor_type != "weight"). + return reference_nvfp4.NVFP4QuantizerRef( + dtype=reference_utils.Fp4Formats.E2M1, + quant_tile_shape=(1, 16), + pow_2_scales=False, + with_rht=with_rht, + row_scaled_nvfp4=row_scaled, + ) if role.tensor_type == "weight": - return quantization_ref_nvfp4.NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + return reference_nvfp4.NVFP4QuantizerRef( + dtype=reference_utils.Fp4Formats.E2M1, quant_tile_shape=(16, 16) if with_2d_quantization else (1, 16), pow_2_scales=False, with_rht=False, @@ -117,6 +143,23 @@ def reset_rng_states(): torch.cuda.manual_seed(seed) +def assert_close_relative(name: str, native, ref, step: int, rel_tol: float) -> None: + """Compare two tensors by relative Frobenius-norm error. + + Used for the row-scaled NVFP4 path, where native (cuBLAS + post-scale) and the + dequantized reference (torch matmul) share the same 4-bit quantization error but + differ at the kernel/impl level (accumulation order, post-scaling). The relative + error therefore stays small even though the absolute NVFP4 error is large. + """ + if native is None or ref is None: + return + ref_f = ref.detach().float() + diff = (native.detach().float() - ref_f).norm().item() + denom = ref_f.norm().item() + rel = diff / denom if denom > 0 else diff + assert rel <= rel_tol, f"{name} relative error {rel:.5f} > {rel_tol} at step {step}" + + def check_nvfp4_module_versus_reference( module_class, in_features: int, @@ -126,6 +169,8 @@ def check_nvfp4_module_versus_reference( num_steps: int = 1, with_rht: bool = False, with_2d_quantization: bool = False, + row_scaled: bool = False, + rel_tol: float = None, ): """ Compare native NVFP4 module against reference implementation. @@ -137,6 +182,10 @@ def check_nvfp4_module_versus_reference( bias: Whether to use bias x_dtype: Input tensor dtype num_steps: Number of forward/backward steps to test + row_scaled: Enable row-scaled (per-token) NVFP4 activation quantization, + exercising the row-scaled columnwise (transpose) path consumed by wgrad. + rel_tol: If set, compare tensors by relative Frobenius-norm error instead of + the default bitwise-tight ``assert_close`` (used for the row-scaled path). """ device = "cuda" batch_size = 32 @@ -203,8 +252,8 @@ def check_nvfp4_module_versus_reference( ref_module.layer_norm_bias.copy_(native_module.layer_norm_bias) # Create recipes for native and reference implementations - nvfp4_recipe = GetRecipes.nvfp4_recipe_to_test(with_rht, with_2d_quantization) - nvfp4_ref_factory = get_nvfp4_quantizer_factory(with_rht, with_2d_quantization) + nvfp4_recipe = GetRecipes.nvfp4_recipe_to_test(with_rht, with_2d_quantization, row_scaled) + nvfp4_ref_factory = get_nvfp4_quantizer_factory(with_rht, with_2d_quantization, row_scaled) nvfp4_ref_recipe = recipe.CustomRecipe(qfactory=nvfp4_ref_factory) # Training loop comparison @@ -279,6 +328,22 @@ def check_nvfp4_module_versus_reference( native_out = native_outputs[step] ref_out = ref_outputs[step] + if rel_tol is not None: + # Row-scaled path: native cuBLAS + post-scale vs dequantized torch + # reference share the same 4-bit error, so compare by relative norm. + assert_close_relative("Output", native_out["output"], ref_out["output"], step, rel_tol) + assert_close_relative( + "Input gradient", native_out["input_grad"], ref_out["input_grad"], step, rel_tol + ) + assert_close_relative( + "Weight gradient", native_out["weight_grad"], ref_out["weight_grad"], step, rel_tol + ) + if bias: + assert_close_relative( + "Bias gradient", native_out["bias_grad"], ref_out["bias_grad"], step, rel_tol + ) + continue + # Compare outputs torch.testing.assert_close( native_out["output"], @@ -367,6 +432,44 @@ def test_nvfp4_linear_versus_reference( ) +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "in_features, out_features", + [ + (128, 256), + (256, 128), + (512, 512), + (768, 3072), + ], +) +@pytest.mark.parametrize("num_steps", [1, 3], ids=["single_step", "multi_step"]) +def test_nvfp4_linear_row_scaled_versus_reference( + in_features: int, + out_features: int, + num_steps: int, +): + """End-to-end row-scaled (per-token) NVFP4 Linear forward + backward. + + Exercises the row-scaled columnwise (transpose) activation quantization that + this PR unblocks: the forward activation is quantized row-scaled with both + rowwise (fprop) and columnwise (wgrad) directions, so ``backward()`` drives the + new transpose path through the wgrad GEMM. Compared against the dequantized + reference quantizer by relative norm (both share the same 4-bit error). + """ + check_nvfp4_module_versus_reference( + module_class=te.Linear, + in_features=in_features, + out_features=out_features, + bias=False, + x_dtype=torch.bfloat16, + num_steps=num_steps, + with_rht=False, + with_2d_quantization=False, + row_scaled=True, + rel_tol=2e-2, + ) + + def check_nvfp4_layernorm_linear_versus_reference( in_features: int, out_features: int, diff --git a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py index 0eebe5d331..b8ceb7ffff 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py @@ -14,8 +14,8 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex from transformer_engine.pytorch import NVFP4Quantizer -from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef -from transformer_engine.pytorch.custom_recipes import utils +from transformer_engine.pytorch.custom_recipes.reference_nvfp4 import NVFP4QuantizerRef +from transformer_engine.pytorch.custom_recipes import reference_utils from transformer_engine.common.recipe import NVFP4BlockScaling recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) @@ -89,7 +89,12 @@ def maybe_skip_row_scaled_unsupported_quantization( if not row_scaled_nvfp4: return if return_transpose: - pytest.skip("Row-scaled NVFP4 does not support columnwise usage") + if use_4over6: + pytest.skip("Row-scaled NVFP4 transpose does not support 4over6 mode") + if x_dtype != torch.bfloat16 or M is None or N is None or M % 32 != 0 or N % 32 != 0: + pytest.skip( + "Row-scaled NVFP4 transpose requires BF16 input and dimensions divisible by 32" + ) if with_2d_quantization: pytest.skip("Row-scaled NVFP4 does not support 2D quantization") @@ -184,7 +189,7 @@ def check_quantization_nvfp4_versus_reference( # Reference quantization quant_tile_shape = (1, 16) if not with_2d_quantization else (16, 16) ref_quantizer = NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + dtype=reference_utils.Fp4Formats.E2M1, rowwise=True, columnwise=return_transpose, pow_2_scales=False, @@ -395,7 +400,7 @@ def test_nvfp4_quantization_extrema_versus_reference( qx_amax_t = x_nvfp4_sut._amax_columnwise ref_quantizer = NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + dtype=reference_utils.Fp4Formats.E2M1, rowwise=True, columnwise=return_transpose, pow_2_scales=False, @@ -542,7 +547,7 @@ def test_nvfp4_quantization_boundary_values( qx_amax_t = x_nvfp4_sut._amax_columnwise ref_quantizer = NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + dtype=reference_utils.Fp4Formats.E2M1, rowwise=True, columnwise=return_transpose, pow_2_scales=False, @@ -675,7 +680,7 @@ def test_nvfp4_quantization_noncontiguous_inputs( qx_amax_t = x_nvfp4_sut._amax_columnwise ref_quantizer = NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + dtype=reference_utils.Fp4Formats.E2M1, rowwise=True, columnwise=return_transpose, pow_2_scales=False, diff --git a/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py index cd9ffe282b..22023bf286 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py @@ -14,8 +14,8 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex from transformer_engine.pytorch import NVFP4Quantizer -from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef -from transformer_engine.pytorch.custom_recipes import utils +from transformer_engine.pytorch.custom_recipes.reference_nvfp4 import NVFP4QuantizerRef +from transformer_engine.pytorch.custom_recipes import reference_utils from transformer_engine.common.recipe import NVFP4BlockScaling from torch.utils.cpp_extension import IS_HIP_EXTENSION @@ -101,7 +101,7 @@ def check_quantization_nvfp4_versus_reference( # Reference quantization using NVFP4QuantizerRef with built-in RHT ref_quantizer = NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + dtype=reference_utils.Fp4Formats.E2M1, rowwise=return_rowwise, columnwise=return_transpose, pow_2_scales=False, diff --git a/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py index 11777a7151..327e70df27 100755 --- a/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py +++ b/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py @@ -444,3 +444,39 @@ def test_group_stochastic_rounding_quantization_versus_reference( num_splits=num_splits, use_tex_split_quantize=use_tex_split_quantize, ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("x_dtype", [torch.float32, torch.bfloat16], ids=str) +@pytest.mark.parametrize("use_2D", [False, True], ids=str) +def test_stochastic_rounding_picks_an_adjacent_fp4_value( + x_dtype: torch.dtype, use_2D: bool +) -> None: + device = "cuda" + torch.manual_seed(seed) + M, N = 512, 1024 + x = torch.randn((M, N), dtype=x_dtype, device=device) + amax = torch.max(torch.abs(x)).float() + + q, s, q_t, s_t = quantize_fp4(x, use_stochastic_rounding=True, use_2D=use_2D, use_RHT=False) + q_redraw, _, q_t_redraw, _ = quantize_fp4( + x, use_stochastic_rounding=True, use_2D=use_2D, use_RHT=False + ) + assert not torch.equal(q, q_redraw), "two stochastic rounding draws are identical" + assert not torch.equal(q_t, q_t_redraw), "two stochastic rounding draws are identical" + + def check_adjacent(qx: torch.Tensor, sx: torch.Tensor, reference: torch.Tensor) -> None: + # The kernel rounds reference / block_scale onto the E2M1 grid, so every + # output has to be one of the two grid values that bracket it, i.e. within + # one grid step. An all-zero output fails this everywhere the input is not + # already near zero. + block_scale = sx.repeat_interleave(16, dim=1).view(torch.float8_e4m3fn).to(torch.float32) + block_scale = block_scale[: reference.shape[0], : reference.shape[1]] * (amax / (6.0 * 448)) + exact = (reference.float() / block_scale).clamp(-6.0, 6.0) + magnitude = exact.abs() + step = torch.where(magnitude >= 4.0, 2.0, torch.where(magnitude >= 2.0, 1.0, 0.5)) + gap = (fp4_to_fp32(unpack_fp4(qx)) - exact).abs() + assert torch.all(gap <= step * 1.0001), f"largest gap to the exact value is {gap.max()}" + + check_adjacent(q, s, x) + check_adjacent(q_t, s_t, x.t().contiguous()) diff --git a/tests/pytorch/test_cpu_offloading.py b/tests/pytorch/test_cpu_offloading.py index ab64534627..e04bd576b5 100644 --- a/tests/pytorch/test_cpu_offloading.py +++ b/tests/pytorch/test_cpu_offloading.py @@ -21,6 +21,10 @@ from transformer_engine.pytorch.fp8 import FP8GlobalStateManager import transformer_engine.pytorch as te from transformer_engine.common import recipe +from hybrid_quantization_utils import ( + hybrid_fp8_mxfp8_qfactory, + hybrid_mxfp8_nvfp4_qfactory, +) from utils import ModelConfig, recipe_id, skip_unsupported_backward_override from torch.utils.cpp_extension import IS_HIP_EXTENSION @@ -69,6 +73,10 @@ def nvfp4_4over6(): quantization_recipes.append(recipe.NVFP4BlockScaling()) quantization_recipes.append(nvfp4_4over6()) quantization_recipes.append(nvfp4_row_scaled()) +if fp8_available and mxfp8_available: + quantization_recipes.append(recipe.CustomRecipe(qfactory=hybrid_fp8_mxfp8_qfactory)) +if mxfp8_available and nvfp4_available: + quantization_recipes.append(recipe.CustomRecipe(qfactory=hybrid_mxfp8_nvfp4_qfactory)) model_config = { @@ -223,6 +231,17 @@ def create_tensor(recipe: Optional[recipe.Recipe], requires_grad: bool = False) nvfp4_use_4over6=use_4over6, ) return quantizer(tensor) + elif recipe.custom(): + # CustomRecipe: invoke the qfactory for the linear weight role + # as a representative quantizer (returns a HybridQuantizer for the + # hybrid factories registered at module scope). + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = recipe.qfactory(QuantizerRole(module_type="linear", tensor_type="weight")) + if quantizer is None: + # Fallback: factory did not supply a weight quantizer. + return tensor.requires_grad_() if requires_grad else tensor + return quantizer(tensor) @staticmethod def create_recipe_ctx(recipe: Optional[recipe.Recipe]): @@ -496,6 +515,16 @@ def test_sanity(self, layer_type, recipe, backward_override): and recipe.float8_block_scaling() ): pytest.skip("Fusible operations do not support FP8 block scaling recipe") + # Skip hybrid (CustomRecipe) on ops-based LayerNormMLP: the ops-based + # LayerNorm passes the quantizer directly to the fused C++ kernel which + # does not recognize HybridQuantizer (cf. design-doc TODO; the regular + # layernorm_mlp.py has an unfused fallback but the ops path does not + # yet). Unrelated to CPU offload. + # grouped_linear is NOT skipped here — it passes test_sanity with + # hybrid; only memory-accounting assertions trip it in test_memory / + # test_manual_synchronization. + if layer_type in ("layernorm_mlp_ops",) and recipe is not None and recipe.custom(): + pytest.skip(f"Hybrid CustomRecipe + {layer_type} integration is not yet complete") recipe_ctx = Utils.create_recipe_ctx(recipe) init_cuda_memory = Utils.get_cuda_memory_mb() @@ -544,6 +573,17 @@ def test_memory(self, layer_type, recipe, backward_override): and recipe.float8_block_scaling() ): pytest.skip("Fusible operations do not support FP8 block scaling recipe") + # Memory-accounting checks fail for grouped_linear with hybrid because + # `_hybrid_split_quantize` produces per-group HybridQuantizedTensorStorage + # whose individual sub-buffers don't all cross the 256K-element offload + # threshold — the net GPU memory drop after offload is smaller than the + # analytical estimate. Correctness (test_sanity, test_numerics) passes. + if ( + layer_type in ("layernorm_mlp_ops", "grouped_linear") + and recipe is not None + and recipe.custom() + ): + pytest.skip(f"Hybrid CustomRecipe + {layer_type} integration is not yet complete") offload_ctx, sync_function = get_cpu_offload_context( enabled=True, @@ -637,6 +677,13 @@ def test_manual_synchronization(self, recipe, layer_type, backward_override): and recipe.float8_block_scaling() ): pytest.skip("Fusible operations do not support FP8 block scaling recipe") + # Same memory-accounting caveat as test_memory (see comment there). + if ( + layer_type in ("layernorm_mlp_ops", "grouped_linear") + and recipe is not None + and recipe.custom() + ): + pytest.skip(f"Hybrid CustomRecipe + {layer_type} integration is not yet complete") offload_ctx, sync_function, manual_controller = get_cpu_offload_context( enabled=True, @@ -716,6 +763,8 @@ def test_numerics( and recipe.float8_block_scaling() ): pytest.skip("Fusible operations do not support FP8 block scaling recipe") + if layer_type in ("layernorm_mlp_ops",) and recipe is not None and recipe.custom(): + pytest.skip(f"Hybrid CustomRecipe + {layer_type} integration is not yet complete") recipe_ctx = Utils.create_recipe_ctx(recipe) @@ -785,71 +834,101 @@ def forward(self, x): ): param_offload.data.copy_(param_no_offload.data) - x = Utils.create_tensor(None) + x = Utils.create_tensor(None, requires_grad=True) if use_cuda_graphs: callable_offload = te.make_graphed_callables( callable_offload, (x,), enabled=recipe is not None, - recipe=(Utils.create_recipe_ctx(recipe) if recipe is not None else None), + recipe=recipe, ) # warm up (for example to compute sf for delayed scaling) for _ in range(4): + callable_offload.zero_grad(set_to_none=True) + callable_no_offload.zero_grad(set_to_none=True) + x.grad = None out = callable_offload(x) out.sum().backward() + x.grad = None out = callable_no_offload(x) out.sum().backward() callable_offload.zero_grad(set_to_none=True) + callable_no_offload.zero_grad(set_to_none=True) + x.grad = None + rng_state = torch.cuda.get_rng_state() out_offload = callable_offload(x) out_offload.sum().backward() - # save out and gradients - offload_outs = [out_offload] + # Clone before the no-offload pass. CUDA graphs may reuse static output + # buffers, and the comparison must cover computed values rather than + # aliases or unchanged parameters. + offload_out = out_offload.detach().clone() + assert x.grad is not None + offload_input_grad = x.grad.detach().clone() + offload_param_grads = [] for param in callable_offload.parameters(): - offload_outs.append(param.detach().clone()) + assert param.grad is not None + offload_param_grads.append(param.grad.detach().clone()) torch.cuda.reset_peak_memory_stats() + torch.cuda.set_rng_state(rng_state) + x.grad = None out_no_offload = callable_no_offload(x) out_no_offload.sum().backward() - # collect gradients - no_offload_outs = [out_no_offload] + no_offload_out = out_no_offload.detach().clone() + assert x.grad is not None + no_offload_input_grad = x.grad.detach().clone() + no_offload_param_grads = [] for param in callable_no_offload.parameters(): - no_offload_outs.append(param.detach().clone()) + assert param.grad is not None + no_offload_param_grads.append(param.grad.detach().clone()) - # check if tensors are the same + # CPU offload is byte-preserving transport. It must not alter forward + # results, input gradients, or any parameter gradient. # - # ROCm (IFU v2.18): the bf16 forward output (tensor 0) can diverge from - # the no-offload path for UnfusedAttention + cuda_graphs + - # Float8CurrentScaling, and ONLY for that config. It reproduces only under - # the full test matrix, never in isolation: a prior non-graphed case leaves - # GPU memory resident, so when make_graphed_callables captures this case's - # graph its private mempool lands over a different allocator state and - # hipBLASLt selects a different (but equally valid) GEMM algorithm at - # capture time. The two algorithms accumulate in a different order, so the - # bf16 output rounds differently. This is GEMM-algorithm nondeterminism, - # not a correctness regression. + # ROCm (IFU v2.18): the bf16 forward output can diverge from the no-offload + # path for UnfusedAttention + cuda_graphs + Float8CurrentScaling, and ONLY + # for that config. It reproduces only under the full test matrix, never in + # isolation: a prior non-graphed case leaves GPU memory resident, so when + # make_graphed_callables captures this case's graph its private mempool lands + # over a different allocator state and hipBLASLt selects a different (but + # equally valid) GEMM algorithm at capture time. The two algorithms accumulate + # in a different order, so the bf16 output rounds differently. This is + # GEMM-algorithm nondeterminism, not a correctness regression. # # Measured worst case across 6 full-matrix runs: max |a-b| = 0.09375 - # (1.5 bf16 ULP at O(1)); the weight and gradient tensors (i>0) stay - # bit-identical in every case/run. So the relaxed tolerance is scoped as - # tightly as the divergence: ROCm only, tensor 0 only, atol=1.5e-1 - # (~1.6x over the 0.09375 worst case). Everything else — all i>0, and CUDA - # for every i — keeps the exact default comparison, so a real weight/grad - # corruption or a CUDA offload regression still fails. + # (1.5 bf16 ULP at O(1)); the input and parameter gradients stay bit-identical + # in every case/run. So the relaxed tolerance is scoped as tightly as the + # divergence: ROCm only, forward output only, atol=1.5e-1 (~1.6x over the + # 0.09375 worst case). The gradients keep the exact (rtol=0, atol=0) + # comparison on both ROCm and CUDA, so a real grad corruption or a CUDA + # offload regression still fails. # See the IFU v2.18 handoff notes for the full bisection + root cause. - for i in range(len(offload_outs)): - if IS_HIP_EXTENSION and i == 0: - assert torch.allclose( - offload_outs[i], no_offload_outs[i], rtol=2e-2, atol=1.5e-1 - ), f"Error in tensor {i}." - else: - assert torch.allclose( - offload_outs[i], no_offload_outs[i] - ), f"Error in tensor {i}." + if IS_HIP_EXTENSION: + torch.testing.assert_close(offload_out, no_offload_out, rtol=2e-2, atol=1.5e-1) + else: + torch.testing.assert_close(offload_out, no_offload_out, rtol=0.0, atol=0.0) + torch.testing.assert_close( + offload_input_grad, + no_offload_input_grad, + rtol=0.0, + atol=0.0, + ) + assert len(offload_param_grads) == len(no_offload_param_grads) + for index, (offload_grad, no_offload_grad) in enumerate( + zip(offload_param_grads, no_offload_param_grads) + ): + torch.testing.assert_close( + offload_grad, + no_offload_grad, + rtol=0.0, + atol=0.0, + msg=f"Parameter gradient {index} differs with CPU offload", + ) torch.cuda.synchronize() diff --git a/tests/pytorch/test_cpu_offloading_v1.py b/tests/pytorch/test_cpu_offloading_v1.py index bcbc91f3b9..7c444486d6 100644 --- a/tests/pytorch/test_cpu_offloading_v1.py +++ b/tests/pytorch/test_cpu_offloading_v1.py @@ -14,6 +14,10 @@ import transformer_engine.pytorch as te from transformer_engine.common import recipe +from hybrid_quantization_utils import ( + hybrid_fp8_mxfp8_qfactory, + hybrid_mxfp8_nvfp4_qfactory, +) from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported from utils import ModelConfig, get_available_attention_backends @@ -21,10 +25,20 @@ # Check supported quantization schemes fp8_available = te.is_fp8_available() mxfp8_available = te.is_mxfp8_available() +nvfp4_available = te.is_nvfp4_available() + quantization_recipes: Optional[recipe.Recipe] = [None] if fp8_available: quantization_recipes.extend((recipe.Float8CurrentScaling(), recipe.DelayedScaling())) +if fp8_available and mxfp8_available: + quantization_recipes.append(recipe.CustomRecipe(qfactory=hybrid_fp8_mxfp8_qfactory)) +if mxfp8_available and nvfp4_available: + quantization_recipes.append(recipe.CustomRecipe(qfactory=hybrid_mxfp8_nvfp4_qfactory)) + +hybrid_quantization_recipes = [ + item for item in quantization_recipes if item is not None and item.custom() +] model_config = { "small": ModelConfig(8, 512, 8, 64, num_layers=5, eps=0.1), @@ -102,6 +116,15 @@ def _estimate_cached_weight_size( if quantization_recipe is None: return 0 + # Hybrid (CustomRecipe) caches two sub-storages per weight with + # potentially different formats. Returning ``None`` signals the caller + # to skip the exact memory-accounting assertion — the ``memory_with_offload + # < memory_without_offload`` check still applies. Deriving an analytical + # estimate here is blocked on the FSDP2-style packing optimization still + # being a TODO in hybrid_quantization_design.md. + if quantization_recipe.custom(): + return None + # Count number of weight param elements param_elements = 0 for module in modules: @@ -186,6 +209,19 @@ def _measure_cached_memory( def test_cpu_offload(quantization_recipe: Optional[recipe.Recipe], model_name: str) -> None: """Check that CPU offloading runs and has expected memory usage.""" + # Skip hybrid (CustomRecipe) on module types whose integration with hybrid + # is not yet complete (preexisting, independent of CPU offload): + # - layernorm_mlp_ops: the ops-based LayerNorm passes the quantizer + # directly to the fused C++ kernel which does not recognize + # HybridQuantizer (cf. design doc; the regular layernorm_mlp.py has + # an unfused fallback but the ops path does not yet). + if ( + model_name in ("layernorm_mlp_ops",) + and quantization_recipe is not None + and quantization_recipe.custom() + ): + pytest.skip(f"Hybrid CustomRecipe + {model_name} integration is not yet complete") + # Construct model modules_list = [model_types[model_name]() for _ in range(NUM_LAYERS)] if model_name in ["multihead_attention", "transformer_layer"]: @@ -214,4 +250,98 @@ def test_cpu_offload(quantization_recipe: Optional[recipe.Recipe], model_name: s modules_list, quantization_recipe, ) - assert abs(memory_with_offload - memory_from_cached_weights) < EPSILON + # ``_estimate_cached_weight_size`` returns ``None`` for recipes whose + # analytical cached-weight size is not worked out (CustomRecipe / hybrid); + # in that case the memory-savings assertion above is the only check. + if memory_from_cached_weights is not None: + assert abs(memory_with_offload - memory_from_cached_weights) < EPSILON + + +@pytest.mark.parametrize( + "quantization_recipe", + hybrid_quantization_recipes, + ids=[item.qfactory.__name__ for item in hybrid_quantization_recipes], +) +def test_hybrid_cpu_offload_numerics_exact(quantization_recipe: recipe.Recipe) -> None: + """V1 offload must preserve hybrid forward and backward values bitwise.""" + + def run(modules, inp, grad_output, *, cpu_offload): + if cpu_offload: + offload_context, sync_function = te.get_cpu_offload_context( + enabled=True, + num_layers=len(modules), + model_layers=len(modules) + 1, + offload_activations=True, + offload_weights=False, + ) + else: + offload_context = contextlib.nullcontext() + sync_function = lambda tensor: tensor + + out = inp + for module in modules: + with te.autocast(enabled=True, recipe=quantization_recipe), offload_context: + out = module(out) + out = sync_function(out) + # Commit the final offload group, mirroring the memory test above. + with offload_context: + out = out.clone() + out = sync_function(out) + out.backward(grad_output) + + grads = {} + for module_index, module in enumerate(modules): + for name, param in module.named_parameters(): + assert param.grad is not None + grads[f"{module_index}.{name}"] = param.grad.detach().clone() + assert inp.grad is not None + return out.detach().clone(), inp.grad.detach().clone(), grads + + torch.manual_seed(9000) + reference_modules = [model_types["linear"]() for _ in range(2)] + torch.manual_seed(9001) + offload_modules = [model_types["linear"]() for _ in range(2)] + with torch.no_grad(): + for offload_module, reference_module in zip(offload_modules, reference_modules): + for offload_param, reference_param in zip( + offload_module.parameters(), reference_module.parameters() + ): + offload_param.copy_(reference_param) + + # V1 requires one pass to initialize cached quantized weight workspaces. + warmup_rng_state = torch.cuda.get_rng_state() + _warmup_model(offload_modules, quantization_recipe) + for module in offload_modules: + module.zero_grad(set_to_none=True) + torch.cuda.set_rng_state(warmup_rng_state) + _warmup_model(reference_modules, quantization_recipe) + for module in reference_modules: + module.zero_grad(set_to_none=True) + + torch.manual_seed(9002) + x = torch.randn((8, SIZE, SIZE), device="cuda", dtype=torch.bfloat16) + grad_output = torch.randn_like(x) + x_offload = x.detach().clone().requires_grad_(True) + x_reference = x.detach().clone().requires_grad_(True) + + # Replay identical stochastic-rounding/RHT randomness in both paths. + rng_state = torch.cuda.get_rng_state() + offload_out, offload_input_grad, offload_param_grads = run( + offload_modules, x_offload, grad_output, cpu_offload=True + ) + torch.cuda.set_rng_state(rng_state) + reference_out, reference_input_grad, reference_param_grads = run( + reference_modules, x_reference, grad_output, cpu_offload=False + ) + + torch.testing.assert_close(offload_out, reference_out, rtol=0.0, atol=0.0) + torch.testing.assert_close(offload_input_grad, reference_input_grad, rtol=0.0, atol=0.0) + assert offload_param_grads.keys() == reference_param_grads.keys() + for name, reference_grad in reference_param_grads.items(): + torch.testing.assert_close( + offload_param_grads[name], + reference_grad, + rtol=0.0, + atol=0.0, + msg=f"V1 CPU-offload parameter gradient mismatch for {name}", + ) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index f07d8b4d9b..5a848dc0e8 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -751,12 +751,296 @@ def test_make_graphed_callables_with_kwargs( assert_all_equal(outputs, graph_outputs) +def test_make_graphed_callables_returns_owned_parameter_grads() -> None: + """Parameter grads returned from graph replay must not alias static graph buffers.""" + reset_rng_states() + model_config = model_configs["small"] + dtype = torch.float32 + model = torch.nn.Linear( + model_config.hidden_size, + model_config.hidden_size, + bias=False, + device="cuda", + dtype=dtype, + ) + model = make_graphed_callables( + model, + (generate_data(model_config, dtype, warmup=True, requires_grad=False),), + ) + + seen_grads = [] + + def save_grad(grad): + seen_grads.append(grad) + return grad + + hook = model.weight.register_hook(save_grad) + try: + output = model(generate_data(model_config, dtype, requires_grad=False)) + output.backward(generate_data(model_config, dtype, requires_grad=False)) + + assert len(seen_grads) == 1 + first_grad = seen_grads[0] + first_grad_ptr = first_grad.data_ptr() + first_grad_snapshot = first_grad.clone() + + model.zero_grad(set_to_none=True) + + output = model(generate_data(model_config, dtype, requires_grad=False)) + output.backward(generate_data(model_config, dtype, requires_grad=False)) + + assert len(seen_grads) == 2 + assert first_grad.data_ptr() == first_grad_ptr + assert seen_grads[1].data_ptr() != first_grad_ptr + torch.testing.assert_close(first_grad, first_grad_snapshot, rtol=0, atol=0) + finally: + hook.remove() + reset_graphs(model) + + +def test_make_graphed_callables_accumulates_owned_parameter_grads() -> None: + """Parameter grad accumulation must not reuse overwritten static graph buffers.""" + reset_rng_states() + model_config = model_configs["small"] + dtype = torch.float32 + model = torch.nn.Linear( + model_config.hidden_size, + model_config.hidden_size, + bias=False, + device="cuda", + dtype=dtype, + ) + model = make_graphed_callables( + model, + (generate_data(model_config, dtype, warmup=True, requires_grad=False),), + ) + + input_1 = generate_data(model_config, dtype, requires_grad=False) + grad_1 = generate_data(model_config, dtype, requires_grad=False) + input_2 = generate_data(model_config, dtype, requires_grad=False) + grad_2 = generate_data(model_config, dtype, requires_grad=False) + expected_grad = torch.einsum("...o,...i->oi", grad_1, input_1) + torch.einsum( + "...o,...i->oi", grad_2, input_2 + ) + + try: + model.zero_grad(set_to_none=True) + model(input_1).backward(grad_1) + model(input_2).backward(grad_2) + torch.testing.assert_close(model.weight.grad, expected_grad, rtol=0, atol=0) + finally: + reset_graphs(model) + + +def test_make_graphed_callables_preserves_skipped_parameter_grad_alias() -> None: + """Delayed-wgrad parameters are excluded from returned-grad clone handling.""" + reset_rng_states() + model_config = model_configs["small"] + dtype = torch.float32 + model = torch.nn.Linear( + model_config.hidden_size, + model_config.hidden_size, + bias=False, + device="cuda", + dtype=dtype, + ) + model.weight.skip_backward_post_hook = True + model = make_graphed_callables( + model, + (generate_data(model_config, dtype, warmup=True, requires_grad=False),), + ) + + seen_grads = [] + + def save_grad(grad): + seen_grads.append(grad) + return grad + + hook = model.weight.register_hook(save_grad) + try: + output = model(generate_data(model_config, dtype, requires_grad=False)) + output.backward(generate_data(model_config, dtype, requires_grad=False)) + + assert len(seen_grads) == 1 + first_grad_ptr = seen_grads[0].data_ptr() + + model.zero_grad(set_to_none=True) + + output = model(generate_data(model_config, dtype, requires_grad=False)) + output.backward(generate_data(model_config, dtype, requires_grad=False)) + + assert len(seen_grads) == 2 + assert seen_grads[1].data_ptr() == first_grad_ptr + finally: + hook.remove() + reset_graphs(model) + + +def test_make_graphed_callables_can_skip_returned_parameter_grad_clone() -> None: + """Parameter grad clone handling can be disabled for callers that manage lifetimes.""" + reset_rng_states() + model_config = model_configs["small"] + dtype = torch.float32 + model = torch.nn.Linear( + model_config.hidden_size, + model_config.hidden_size, + bias=False, + device="cuda", + dtype=dtype, + ) + model = make_graphed_callables( + model, + (generate_data(model_config, dtype, warmup=True, requires_grad=False),), + clone_param_grads_on_return=False, + ) + + seen_grads = [] + + def save_grad(grad): + seen_grads.append(grad) + return grad + + hook = model.weight.register_hook(save_grad) + try: + output = model(generate_data(model_config, dtype, requires_grad=False)) + output.backward(generate_data(model_config, dtype, requires_grad=False)) + + assert len(seen_grads) == 1 + first_grad_ptr = seen_grads[0].data_ptr() + + model.zero_grad(set_to_none=True) + + output = model(generate_data(model_config, dtype, requires_grad=False)) + output.backward(generate_data(model_config, dtype, requires_grad=False)) + + assert len(seen_grads) == 2 + assert seen_grads[1].data_ptr() == first_grad_ptr + finally: + hook.remove() + reset_graphs(model) + + +def test_make_graphed_callables_snapshots_parameter_grad_clone_policy() -> None: + """Parameter grad clone policy is fixed at capture time.""" + reset_rng_states() + model_config = model_configs["small"] + dtype = torch.float32 + model = torch.nn.Linear( + model_config.hidden_size, + model_config.hidden_size, + bias=False, + device="cuda", + dtype=dtype, + ) + model = make_graphed_callables( + model, + (generate_data(model_config, dtype, warmup=True, requires_grad=False),), + ) + model.weight.skip_backward_post_hook = True + + seen_grads = [] + + def save_grad(grad): + seen_grads.append(grad) + return grad + + hook = model.weight.register_hook(save_grad) + try: + output = model(generate_data(model_config, dtype, requires_grad=False)) + output.backward(generate_data(model_config, dtype, requires_grad=False)) + + assert len(seen_grads) == 1 + first_grad = seen_grads[0] + first_grad_ptr = first_grad.data_ptr() + first_grad_snapshot = first_grad.clone() + + model.zero_grad(set_to_none=True) + + output = model(generate_data(model_config, dtype, requires_grad=False)) + output.backward(generate_data(model_config, dtype, requires_grad=False)) + + assert len(seen_grads) == 2 + assert seen_grads[1].data_ptr() != first_grad_ptr + torch.testing.assert_close(first_grad, first_grad_snapshot, rtol=0, atol=0) + finally: + hook.remove() + reset_graphs(model) + + +def _make_capture_time_hooks( + modules: Tuple[torch.nn.Module, ...], + records: List[Tuple[int, str]], +) -> List[Dict[str, Dict[int, Callable]]]: + """Make capture-time hooks that record call order.""" + + def make_hook(module_idx: int, hook_name: str) -> Callable: + expected_module = modules[module_idx] + + def hook(module: torch.nn.Module) -> None: + assert module is expected_module + assert not torch.cuda.is_current_stream_capturing() + records.append((module_idx, hook_name)) + + return hook + + return [ + { + "forward_pre_hooks": {0: make_hook(module_idx, "forward_pre_hooks")}, + "forward_hooks": {0: make_hook(module_idx, "forward_hooks")}, + "backward_pre_hooks": {0: make_hook(module_idx, "backward_pre_hooks")}, + "backward_hooks": {0: make_hook(module_idx, "backward_hooks")}, + } + for module_idx in range(len(modules)) + ] + + +@pytest.mark.parametrize("with_order", (False, True)) +def test_make_graphed_callables_with_capture_time_hooks(with_order: bool) -> None: + """Test capture-time hooks around warmup and graph capture.""" + num_warmup_iters = 2 + modules = ( + torch.nn.Linear(8, 8, device="cuda"), + torch.nn.Linear(8, 8, device="cuda"), + ) + sample_args = tuple((torch.ones(4, 8, device="cuda", requires_grad=True),) for _ in modules) + records = [] + hook_order = [ + (0, "forward_pre_hooks"), + (0, "forward_hooks"), + (1, "forward_pre_hooks"), + (1, "forward_hooks"), + (1, "backward_pre_hooks"), + (1, "backward_hooks"), + (0, "backward_pre_hooks"), + (0, "backward_hooks"), + ] + + graphed_callables = make_graphed_callables( + modules, + sample_args, + num_warmup_iters=num_warmup_iters, + _order=[1, 2, -2, -1] if with_order else None, + capture_time_hooks=_make_capture_time_hooks(modules, records), + ) + + assert records == hook_order * (num_warmup_iters + 1) + + for graphed in graphed_callables: + x = torch.randn(4, 8, device="cuda", requires_grad=True) + y = graphed(x) + y.backward(torch.ones_like(y)) + assert records == hook_order * (num_warmup_iters + 1) + reset_graphs(graphed_callables) + + def _test_cuda_graphs_with_interleaved_pipeline_parallelism( *, with_graph: bool, model_config: ModelConfig, dtype: torch.dtype, -) -> List[torch.Tensor]: + reuse_graph_input_output_buffers: bool = False, + clone_param_grads_on_return: bool = True, +) -> Tuple[List[torch.Tensor], List[torch.Tensor]]: """Simulate Megatron-LM interleaved pipeline parallelism.""" reset_rng_states() @@ -796,6 +1080,8 @@ def _test_cuda_graphs_with_interleaved_pipeline_parallelism( sample_args, allow_unused_input=True, _order=layer_order, + _reuse_graph_input_output_buffers=reuse_graph_input_output_buffers, + clone_param_grads_on_return=clone_param_grads_on_return, ) layer_forwards = { (i // num_microbatches, i % num_microbatches): forward @@ -822,11 +1108,15 @@ def _test_cuda_graphs_with_interleaved_pipeline_parallelism( # Cache for layer outputs. outputs = {} + output_snapshots = {} if reuse_graph_input_output_buffers else None def forward(layer_idx: int, microbatch_idx: int): """Helper function for forward steps""" idxs = (layer_idx, microbatch_idx) outputs[idxs] = layer_forwards[idxs](inputs[idxs]) + if output_snapshots is not None: + # Reused graph output buffers are only valid until their corresponding backward. + output_snapshots[idxs] = outputs[idxs].detach().clone() def backward(layer_idx: int, microbatch_idx: int): """Helper function for backward steps""" @@ -849,11 +1139,13 @@ def backward(layer_idx: int, microbatch_idx: int): # Optimizer step. optimizer.step() - outputs = [y for _, y in sorted(outputs.items())] - outputs = get_outputs(model, outputs) + output_values = output_snapshots if output_snapshots is not None else outputs + output_values = [y for _, y in sorted(output_values.items())] + outputs = get_outputs(model, output_values) + final_weights = [param.detach().clone() for param in model.parameters()] if with_graph: reset_graphs(layer_forwards) - return outputs + return outputs, final_weights def test_make_graphed_callables_with_interleaved_pipeline_parallelism( @@ -864,12 +1156,56 @@ def test_make_graphed_callables_with_interleaved_pipeline_parallelism( """Test CUDA graphs with Megatron-LM interleaved pipeline parallelism.""" model_config = model_configs[model_config] kwargs = dict(model_config=model_config, dtype=dtype) - outputs = _test_cuda_graphs_with_interleaved_pipeline_parallelism( + outputs, weights = _test_cuda_graphs_with_interleaved_pipeline_parallelism( + with_graph=False, + **kwargs, + ) + graph_outputs, graph_weights = _test_cuda_graphs_with_interleaved_pipeline_parallelism( + with_graph=True, + **kwargs, + ) + assert_all_equal(outputs, graph_outputs) + assert_all_equal(weights, graph_weights) + + +def test_make_graphed_callables_with_interleaved_pipeline_parallelism_reused_buffers( + *, + model_config: str = "small", + dtype: torch.dtype = torch.float16, +) -> None: + """Test CUDA graphs with reused input/output buffers.""" + model_config = model_configs[model_config] + kwargs = dict(model_config=model_config, dtype=dtype) + outputs, weights = _test_cuda_graphs_with_interleaved_pipeline_parallelism( + with_graph=False, + **kwargs, + ) + graph_outputs, graph_weights = _test_cuda_graphs_with_interleaved_pipeline_parallelism( + with_graph=True, + reuse_graph_input_output_buffers=True, + **kwargs, + ) + assert_all_equal(outputs, graph_outputs) + assert_all_equal(weights, graph_weights) + + +def test_make_graphed_callables_with_interleaved_pipeline_parallelism_reused_buffers_no_param_grad_clone( + *, + model_config: str = "small", + dtype: torch.dtype = torch.float16, +) -> None: + """Test reused input/output buffers when returned parameter grad clones are disabled.""" + model_config = model_configs[model_config] + kwargs = dict(model_config=model_config, dtype=dtype) + outputs, weights = _test_cuda_graphs_with_interleaved_pipeline_parallelism( with_graph=False, **kwargs, ) - graph_outputs = _test_cuda_graphs_with_interleaved_pipeline_parallelism( + graph_outputs, graph_weights = _test_cuda_graphs_with_interleaved_pipeline_parallelism( with_graph=True, + reuse_graph_input_output_buffers=True, + clone_param_grads_on_return=False, **kwargs, ) assert_all_equal(outputs, graph_outputs) + assert_all_equal(weights, graph_weights) diff --git a/tests/pytorch/test_custom_recipe.py b/tests/pytorch/test_custom_recipe.py index 3e6fdb816b..a1a6ef7b37 100644 --- a/tests/pytorch/test_custom_recipe.py +++ b/tests/pytorch/test_custom_recipe.py @@ -7,27 +7,39 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex +from hybrid_quantization_utils import ( + nvfp4_linear_mxfp8_dpa_test_factory as _nvfp4_linear_mxfp8_dpa_factory, +) from transformer_engine.common import recipe from transformer_engine.pytorch.constants import FP8BwdTensorIdx, FP8FwdTensorIdx from transformer_engine.pytorch import ( autocast, + Fp8Padding, + Fp8Unpadding, Linear, LayerNormLinear, LayerNormMLP, GroupedLinear, Float8CurrentScalingQuantizer, ) -from transformer_engine.pytorch.quantization import QuantizerRole +from transformer_engine.pytorch.quantization import ( + QuantizerRole, + get_align_size_for_quantization, +) import transformer_engine.pytorch.ops as te_ops -from transformer_engine.pytorch.custom_recipes.quantization_recipes_base import ( - current_scaling_quantizer_factory, - mxfp8_quantizer_factory, - float8_block_scaling_quantizer_factory, - nvfp4_quantizer_factory, - delayed_scaling_quantizer_factory, +from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + current_scaling_factory, + mxfp8_factory, + float8_block_scaling_factory, + nvfp4_factory, + delayed_scaling_factory, + high_precision_factory, ) -from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import ( - nvfp4_ref_rht_2d_quantizer_factory, +from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + mxfp8_fwd_nvfp4_bwd_factory, +) +from transformer_engine.pytorch.custom_recipes.reference_nvfp4 import ( + nvfp4_ref_rht_2d_factory, ) @@ -58,7 +70,7 @@ def test_custom_recipe_sanity_modules_nvfp4(module_type): inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True) # Use NVFP4 quantizer factory - custom_recipe = recipe.CustomRecipe(qfactory=nvfp4_ref_rht_2d_quantizer_factory) + custom_recipe = recipe.CustomRecipe(qfactory=nvfp4_ref_rht_2d_factory) # Execute with custom recipe with autocast(enabled=True, recipe=custom_recipe): @@ -389,7 +401,7 @@ def _assert_match(out_ref, out_cus, grad_ref, grad_cus, pgrads_ref, pgrads_cus): def test_factory_matches_delayed_scaling(): - """delayed_scaling_quantizer_factory should produce bit-identical results + """delayed_scaling_factory should produce bit-identical results to the built-in DelayedScaling recipe.""" available, reason = te.is_fp8_available(return_reason=True) if not torch.cuda.is_available() or not available: @@ -399,13 +411,13 @@ def test_factory_matches_delayed_scaling(): out_ref, grad_ref, pgrads_ref = _run_linear_fwd_bwd(model_ref, inp_ref, recipe.DelayedScaling()) out_cus, grad_cus, pgrads_cus = _run_linear_fwd_bwd( - model_cus, inp_cus, recipe.CustomRecipe(qfactory=delayed_scaling_quantizer_factory) + model_cus, inp_cus, recipe.CustomRecipe(qfactory=delayed_scaling_factory) ) _assert_match(out_ref, out_cus, grad_ref, grad_cus, pgrads_ref, pgrads_cus) def test_factory_matches_current_scaling(): - """current_scaling_quantizer_factory should produce bit-identical results + """current_scaling_factory should produce bit-identical results to the built-in Float8CurrentScaling recipe.""" available, reason = te.is_fp8_available(return_reason=True) if not torch.cuda.is_available() or not available: @@ -417,13 +429,13 @@ def test_factory_matches_current_scaling(): model_ref, inp_ref, recipe.Float8CurrentScaling() ) out_cus, grad_cus, pgrads_cus = _run_linear_fwd_bwd( - model_cus, inp_cus, recipe.CustomRecipe(qfactory=current_scaling_quantizer_factory) + model_cus, inp_cus, recipe.CustomRecipe(qfactory=current_scaling_factory) ) _assert_match(out_ref, out_cus, grad_ref, grad_cus, pgrads_ref, pgrads_cus) def test_factory_matches_mxfp8(): - """mxfp8_quantizer_factory should produce bit-identical results + """mxfp8_factory should produce bit-identical results to the built-in MXFP8BlockScaling recipe.""" available, reason = te.is_mxfp8_available(return_reason=True) if not torch.cuda.is_available() or not available: @@ -435,13 +447,13 @@ def test_factory_matches_mxfp8(): model_ref, inp_ref, recipe.MXFP8BlockScaling() ) out_cus, grad_cus, pgrads_cus = _run_linear_fwd_bwd( - model_cus, inp_cus, recipe.CustomRecipe(qfactory=mxfp8_quantizer_factory) + model_cus, inp_cus, recipe.CustomRecipe(qfactory=mxfp8_factory) ) _assert_match(out_ref, out_cus, grad_ref, grad_cus, pgrads_ref, pgrads_cus) def test_factory_matches_block_scaling(): - """float8_block_scaling_quantizer_factory should produce bit-identical results + """float8_block_scaling_factory should produce bit-identical results to the built-in Float8BlockScaling recipe.""" available = te.is_fp8_block_scaling_available() if not torch.cuda.is_available() or not available: @@ -453,13 +465,13 @@ def test_factory_matches_block_scaling(): model_ref, inp_ref, recipe.Float8BlockScaling() ) out_cus, grad_cus, pgrads_cus = _run_linear_fwd_bwd( - model_cus, inp_cus, recipe.CustomRecipe(qfactory=float8_block_scaling_quantizer_factory) + model_cus, inp_cus, recipe.CustomRecipe(qfactory=float8_block_scaling_factory) ) _assert_match(out_ref, out_cus, grad_ref, grad_cus, pgrads_ref, pgrads_cus) def test_factory_matches_nvfp4(): - """nvfp4_quantizer_factory should produce bit-identical results + """nvfp4_factory should produce bit-identical results to the built-in NVFP4BlockScaling recipe.""" available = te.is_nvfp4_available() if not torch.cuda.is_available() or not available: @@ -471,7 +483,41 @@ def test_factory_matches_nvfp4(): model_ref, inp_ref, recipe.NVFP4BlockScaling() ) out_cus, grad_cus, pgrads_cus = _run_linear_fwd_bwd( - model_cus, inp_cus, recipe.CustomRecipe(qfactory=nvfp4_quantizer_factory) + model_cus, inp_cus, recipe.CustomRecipe(qfactory=nvfp4_factory) + ) + + _assert_match(out_ref, out_cus, grad_ref, grad_cus, pgrads_ref, pgrads_cus) + + +def test_factory_matches_high_precision(): + """high_precision_factory (all IdentityQuantizer) should produce bit-identical + results to running with no quantization (plain BF16, no autocast). + + Unlike the other ``test_factory_matches_*`` cases, the reference here is not a + built-in recipe but the unquantized path: IdentityQuantizer leaves tensors + untouched, so the GEMMs run in high precision exactly as they would with + autocast disabled. + + Note: the custom side still runs inside ``autocast(enabled=True, ...)``, which + asserts FP8 availability on entry, so this test skips on non-FP8 hardware even + though no tensor is actually quantized. + """ + available, reason = te.is_fp8_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 unsupported: {reason}") + + model_ref, model_cus, inp_ref, inp_cus = _make_pair() + + # Reference: no autocast -> high-precision GEMMs, no quantization. + with autocast(enabled=False): + out_ref = model_ref(inp_ref) + out_ref.float().sum().backward() + out_ref = out_ref.clone() + grad_ref = inp_ref.grad.clone() + pgrads_ref = {n: p.grad.clone() for n, p in model_ref.named_parameters() if p.grad is not None} + + out_cus, grad_cus, pgrads_cus = _run_linear_fwd_bwd( + model_cus, inp_cus, recipe.CustomRecipe(qfactory=high_precision_factory) ) _assert_match(out_ref, out_cus, grad_ref, grad_cus, pgrads_ref, pgrads_cus) @@ -568,30 +614,30 @@ def targeting_factory(role): recorded_roles.append(role) if role is None: - return nvfp4_quantizer_factory(role) + return nvfp4_factory(role) assert isinstance(role, QuantizerRole), f"Expected QuantizerRole, got {type(role)}" # Layer 0 (tl0.*): all MXFP8 if role.name.startswith("tl0"): - return mxfp8_quantizer_factory(role) + return mxfp8_factory(role) # Layer 1 (tl1.*): NVFP4 default, but fc2 overridden to MXFP8 if role.name == "tl1.layernorm_mlp.fc2": - return mxfp8_quantizer_factory(role) + return mxfp8_factory(role) # Layer 2: block scaling for qkv and fc1, rest falls through to default if role.name == "tl2.self_attention.layernorm_linear_qkv": - return float8_block_scaling_quantizer_factory(role) + return float8_block_scaling_factory(role) if role.name == "tl2.layernorm_mlp.fc1": - return float8_block_scaling_quantizer_factory(role) + return float8_block_scaling_factory(role) # Layer 3: current-scaling for proj, rest falls through to default if role.name == "tl3.proj": - return current_scaling_quantizer_factory(role) + return current_scaling_factory(role) # Default: NVFP4 - return nvfp4_quantizer_factory(role) + return nvfp4_factory(role) custom_recipe = recipe.CustomRecipe(qfactory=targeting_factory) @@ -1063,7 +1109,7 @@ def test_role_change_does_not_invalidate_when_role_unchanged(): def test_custom_recipe_dpa_fp8(): - """DotProductAttention forward+backward with CustomRecipe and role-based mixed quantizers. + """DotProductAttention forward+backward with CustomRecipe and role-based DPA quantizers. Uses the nvfp4_linear_fp8_dpa_factory which dispatches: * DPA S/dP slots -> DelayedScalingRequest (stateful) @@ -1090,7 +1136,7 @@ def test_custom_recipe_dpa_fp8(): Float8Quantizer, Float8CurrentScalingQuantizer, ) - from transformer_engine.pytorch.custom_recipes.quantization_factory_examples import ( + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( nvfp4_linear_fp8_dpa_factory, ) @@ -1190,15 +1236,15 @@ def test_custom_recipe_dpa_fp8(): def test_custom_recipe_dpa_mxfp8(): """DotProductAttention forward+backward with CustomRecipe and MXFP8 attention. - Uses the nvfp4_linear_mxfp8_dpa_factory which dispatches: + Uses a test-local reference factory which dispatches: * DPA roles (QKV/O/S/dO/dP/dQKV) -> MXFP8Quantizer (S/dP later nulled out by ``get_attention_quantizers`` since the MXFP8 fused-attention kernel handles those slots internally) * DPA boundary hints -> MXFP8Quantizer * Linear slots -> NVFP4Quantizer - Mirrors the documented "NVFP4 linear + MXFP8 attention" combo from - ``dot_product_attention.py``'s recipe-combination table. + Mirrors the documented experimental "NVFP4 Linear + MXFP8 attention" + configuration. """ available, reason = te.is_fp8_available(return_reason=True) if not torch.cuda.is_available() or not available: @@ -1217,9 +1263,6 @@ def test_custom_recipe_dpa_mxfp8(): from transformer_engine.pytorch.quantization import CustomRecipeState from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer - from transformer_engine.pytorch.custom_recipes.quantization_factory_examples import ( - nvfp4_linear_mxfp8_dpa_factory, - ) torch.manual_seed(42) @@ -1239,7 +1282,7 @@ def test_custom_recipe_dpa_mxfp8(): out_proj = Linear(H, H, params_dtype=torch.bfloat16, bias=False, name="proj").cuda() custom_recipe = recipe.CustomRecipe( - qfactory=nvfp4_linear_mxfp8_dpa_factory, + qfactory=_nvfp4_linear_mxfp8_dpa_factory, fp8_dpa=True, ) @@ -1364,7 +1407,7 @@ def test_custom_recipe_debug_tool_compat(): in_features, out_features, params_dtype=torch.bfloat16, name="layer" ).cuda() - custom_recipe = recipe.CustomRecipe(qfactory=current_scaling_quantizer_factory) + custom_recipe = recipe.CustomRecipe(qfactory=current_scaling_factory) assert TEDebugState.debug_enabled, "Debug mode should be active" @@ -1849,3 +1892,94 @@ def test_slot_role_supports_module_type_only_role(): assert resolved.name == "" # Tensor-type-only recipes fall back to positional for this slot. assert state._slot_tensor_type(0) == "input" + + +_alignment_fp8_available, _alignment_fp8_reason = te.is_fp8_available(return_reason=True) +_alignment_mxfp8_available, _alignment_mxfp8_reason = te.is_mxfp8_available(return_reason=True) +_alignment_nvfp4_available, _alignment_nvfp4_reason = te.is_nvfp4_available(return_reason=True) + + +def test_custom_recipe_quantization_alignment_contract(): + """The alignment contract is safe, configurable, and does not invoke qfactory.""" + calls = [] + + def qfactory(role): + calls.append(role) + return current_scaling_factory(role) + + custom_recipe = recipe.CustomRecipe(qfactory=qfactory) + assert get_align_size_for_quantization(custom_recipe) == 128 + assert calls == [] + + custom_recipe = recipe.CustomRecipe(qfactory=qfactory, quantization_alignment=32) + assert get_align_size_for_quantization(custom_recipe) == 32 + assert calls == [] + + with pytest.raises(ValueError, match="quantization_alignment must be positive"): + recipe.CustomRecipe(qfactory=qfactory, quantization_alignment=0) + + +@pytest.mark.parametrize( + "qfactory", + [ + pytest.param( + current_scaling_factory, + marks=pytest.mark.skipif( + not _alignment_fp8_available, + reason=_alignment_fp8_reason, + ), + id="fp8_current_scaling", + ), + pytest.param( + mxfp8_factory, + marks=pytest.mark.skipif( + not _alignment_mxfp8_available, + reason=_alignment_mxfp8_reason, + ), + id="mxfp8", + ), + pytest.param( + nvfp4_factory, + marks=pytest.mark.skipif( + not _alignment_nvfp4_available, + reason=_alignment_nvfp4_reason, + ), + id="nvfp4", + ), + pytest.param( + mxfp8_fwd_nvfp4_bwd_factory, + marks=pytest.mark.skipif( + not (_alignment_mxfp8_available and _alignment_nvfp4_available), + reason=f"MXFP8: {_alignment_mxfp8_reason}; NVFP4: {_alignment_nvfp4_reason}", + ), + id="hybrid_mxfp8_nvfp4", + ), + ], +) +def test_custom_recipe_automatic_padding_bad_splits(qfactory): + """Automatic alignment makes misaligned CustomRecipe grouped GEMMs valid.""" + custom_recipe = recipe.CustomRecipe(qfactory=qfactory) + padding = Fp8Padding(2) + unpadding = Fp8Unpadding(2) + grouped_linear = GroupedLinear( + 2, + 128, + 128, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + ) + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + original_splits = [17, 15] + + with autocast(enabled=True, recipe=custom_recipe): + padded_inp, padded_splits = padding(inp, original_splits) + assert padding.align_size == get_align_size_for_quantization(custom_recipe) == 128 + assert padded_splits == [128, 128] + padded_out = grouped_linear(padded_inp, padded_splits) + out = unpadding(padded_out, original_splits) + + out.sum().backward() + assert out.shape == (sum(original_splits), 128) + assert torch.isfinite(out).all() + assert torch.isfinite(inp.grad).all() diff --git a/tests/pytorch/test_distributed_weight.py b/tests/pytorch/test_distributed_weight.py new file mode 100644 index 0000000000..69b5917d90 --- /dev/null +++ b/tests/pytorch/test_distributed_weight.py @@ -0,0 +1,172 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Unit tests for the GTP-agnostic DistributedWeight protocol and dispatchers. + +These tests exercise TE's weight-parallelism extension point in isolation, with +a tiny in-repo ``FakeDistributedWeight`` stub standing in for a real implementer +(e.g. Megatron's GTPShardedParam). No GPU, process group, or Megatron import is +required: the whole contract lives in TE and is verified against the fake. +""" + +import pytest +import torch + +from transformer_engine.pytorch.distributed_weight import ( + DistributedWeight, + is_distributed_weight, + materialize_weight_for_forward, + materialize_weight_for_backward, + finalize_weight_grads, +) + + +class FakeDistributedWeight(torch.Tensor): + """Minimal DistributedWeight implementer for dispatcher tests. + + Records how it was called and returns marker tensors so the dispatcher's + behavior (delegation, list normalization, no-op fallback) is observable. + """ + + is_distributed_weight = True + + def __new__(cls, group_size=1): + t = torch.zeros(1).as_subclass(cls) + t.group_size = group_size + t.calls = [] + return t + + def materialize_group_for_forward(self): + self.calls.append("fwd") + out = [torch.full((2, 2), float(i)) for i in range(self.group_size)] + # Match the real GTP contract: single weight returns a bare tensor. + return out if self.group_size > 1 else out[0] + + def materialize_group_for_backward(self, **kwargs): + self.calls.append(("bwd", kwargs)) + out = [torch.full((2, 2), float(10 + i)) for i in range(self.group_size)] + return out if self.group_size > 1 else out[0] + + def finalize_group_grads(self, wgrads, **kwargs): + self.calls.append(("finalize", wgrads)) + wl = wgrads if isinstance(wgrads, (list, tuple)) else [wgrads] + out = [w + 100 for w in wl] + return out if self.group_size > 1 else out[0] + + def grad_buffer(self): + return torch.full((2, 2), -1.0) + + +class FakeNonTensorWeight: + """DistributedWeight-shaped object that is NOT a torch.Tensor (contract violation).""" + + is_distributed_weight = True + + def materialize_group_for_forward(self): + return torch.zeros(2, 2) + + def materialize_group_for_backward(self, **kwargs): + return torch.zeros(2, 2) + + def finalize_group_grads(self, wgrads, **kwargs): + return wgrads + + def grad_buffer(self): + return torch.zeros(2, 2) + + +def test_protocol_runtime_checkable(): + """A conforming object passes isinstance; a plain tensor does not.""" + assert isinstance(FakeDistributedWeight(), DistributedWeight) + assert not isinstance(torch.zeros(2), DistributedWeight) + + +def test_is_distributed_weight(): + assert is_distributed_weight(FakeDistributedWeight()) + assert not is_distributed_weight(torch.zeros(2)) + assert not is_distributed_weight(torch.nn.Parameter(torch.zeros(2))) + + +def test_non_tensor_implementer_rejected(): + """Implementers must be torch.Tensor subclasses; a non-Tensor fails loudly.""" + with pytest.raises(TypeError, match="torch.Tensor subclass"): + is_distributed_weight(FakeNonTensorWeight()) + + +def test_forward_noop_on_plain_tensor(): + """Plain weights pass through unchanged — the critical non-regression.""" + w = torch.nn.Parameter(torch.randn(4, 4)) + out = materialize_weight_for_forward(w) + assert out == [w] + assert out[0] is w + + +@pytest.mark.parametrize("group_size", [1, 3]) +def test_forward_dispatches(group_size): + """Linear (N=1) and GroupedLinear (N=k): one coalesced call, full list returned.""" + w = FakeDistributedWeight(group_size=group_size) + out = materialize_weight_for_forward(w) + assert isinstance(out, list) and len(out) == group_size + # Leader is delegated to exactly once (coalesced), not once per weight. + assert w.calls == ["fwd"] + + +def test_forward_accepts_weight_list(): + """The dispatcher accepts the full per-expert list; the leader (index 0) coalesces it.""" + leader = FakeDistributedWeight(group_size=3) + followers = [torch.zeros(2, 2), torch.zeros(2, 2)] + out = materialize_weight_for_forward([leader, *followers]) + assert isinstance(out, list) and len(out) == 3 + # Leader delegated exactly once; the follower entries are not materialized separately. + assert leader.calls == ["fwd"] + + +def test_forward_noop_on_plain_weight_list(): + """A non-distributed weight list passes through unchanged (all N returned).""" + ws = [torch.nn.Parameter(torch.randn(2, 2)) for _ in range(3)] + out = materialize_weight_for_forward(ws) + assert out == ws + + +@pytest.mark.parametrize("group_size", [1, 2]) +def test_backward_dispatches(group_size): + w = FakeDistributedWeight(group_size=group_size) + out = materialize_weight_for_backward(w) + assert len(out) == group_size + assert torch.equal(out[0], torch.full((2, 2), 10.0)) + + +def test_backward_accepts_weight_list(): + """Backward dispatcher also accepts the full per-expert list; the leader coalesces it.""" + leader = FakeDistributedWeight(group_size=2) + out = materialize_weight_for_backward([leader, torch.zeros(2, 2)]) + assert isinstance(out, list) and len(out) == 2 + assert torch.equal(out[0], torch.full((2, 2), 10.0)) + + +def test_backward_noop_on_plain_tensor(): + plain = torch.zeros(2) + assert materialize_weight_for_backward(plain) == [plain] + + +def test_backward_noop_on_plain_weight_list(): + """A non-distributed weight list passes through unchanged (all N returned).""" + ws = [torch.nn.Parameter(torch.randn(2, 2)) for _ in range(3)] + assert materialize_weight_for_backward(ws) == ws + + +@pytest.mark.parametrize("group_size", [1, 2]) +def test_finalize_grads_dispatches(group_size): + w = FakeDistributedWeight(group_size=group_size) + wgrads = [torch.zeros(2, 2) for _ in range(group_size)] + out = finalize_weight_grads(w, wgrads) + assert len(out) == group_size + assert torch.equal(out[0], torch.full((2, 2), 100.0)) + + +def test_finalize_grads_noop_on_plain_tensor(): + """No-op path leaves the grads untouched.""" + plain_w = torch.nn.Parameter(torch.zeros(2)) + g = [torch.ones(2)] + assert finalize_weight_grads(plain_w, g) == g diff --git a/tests/pytorch/test_float8_current_scaling_exact.py b/tests/pytorch/test_float8_current_scaling_exact.py index f4010d81c3..689c68235c 100644 --- a/tests/pytorch/test_float8_current_scaling_exact.py +++ b/tests/pytorch/test_float8_current_scaling_exact.py @@ -15,8 +15,8 @@ from transformer_engine.common.recipe import Float8CurrentScaling, Format from transformer_engine.pytorch.quantization import autocast, get_fp8_torch_dtype from transformer_engine.pytorch.constants import TE_DType -from transformer_engine.pytorch.custom_recipes.quantization import MMParams -from transformer_engine.pytorch.custom_recipes.quantization_ref_current_scaling import ( +from transformer_engine.pytorch.custom_recipes.gemm import MMParams +from transformer_engine.pytorch.custom_recipes.reference_current_scaling import ( CurrentScalingQuantizerRef, ) from transformer_engine.pytorch.utils import get_torch_float8_e4m3_type diff --git a/tests/pytorch/test_float8blockwisetensor.py b/tests/pytorch/test_float8blockwisetensor.py index a53a174f92..00e111a14d 100644 --- a/tests/pytorch/test_float8blockwisetensor.py +++ b/tests/pytorch/test_float8blockwisetensor.py @@ -102,6 +102,50 @@ def test_constructor( assert tensor.dtype == dtype, "Incorrect nominal dtype" assert tensor.is_cuda, "Incorrect device" + def test_fsdp_assign_gathered_rejects_invalid_scale_geometry(self): + """Reject gathered scales that cannot describe the gathered data tiles.""" + shape = (128, 256) + quantizer = Float8BlockQuantizer( + fp8_dtype=DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=2, + ) + tensor = Float8BlockwiseQTensor( + shape=shape, + dtype=torch.bfloat16, + rowwise_data=torch.zeros(shape, device="cuda", dtype=torch.uint8), + rowwise_scale_inv=torch.ones( + quantizer.get_scale_shape(shape, columnwise=False), + device="cuda", + dtype=torch.float32, + ), + columnwise_data=torch.zeros( + quantizer.get_columnwise_shape(shape), + device="cuda", + dtype=torch.uint8, + ), + columnwise_scale_inv=torch.ones( + quantizer.get_scale_shape(shape, columnwise=True), + device="cuda", + dtype=torch.float32, + ), + fp8_dtype=DType.kFloat8E4M3, + is_2D_scaled=True, + quantizer=quantizer, + ) + original_data = tensor._rowwise_data + gathered_data = torch.zeros((256, 256), device="cuda", dtype=torch.uint8) + invalid_scale = torch.ones((3, 2), device="cuda", dtype=torch.float32) + + with pytest.raises(RuntimeError, match="gathered scale geometry"): + tensor.fsdp_assign_gathered( + (gathered_data, invalid_scale), + {"direction": "rowwise"}, + ) + + assert tensor._rowwise_data is original_data + def _test_quantize_dequantize( self, quantizer: Float8BlockQuantizer, diff --git a/tests/pytorch/test_fused_optimizer.py b/tests/pytorch/test_fused_optimizer.py index 6a9c77ca33..8de200cdf2 100644 --- a/tests/pytorch/test_fused_optimizer.py +++ b/tests/pytorch/test_fused_optimizer.py @@ -171,6 +171,31 @@ def test_frozen_model(self): torch.testing.assert_close(ref_param, tst_param) + @pytest.mark.parametrize("capturable", [False, True]) + def test_empty_param_group_advances_step(self, capturable): + # An empty param group must advance its step counter like a populated one. + # The same group can be empty on one data-parallel rank and populated on + # another, and "step" is part of the checkpoint, so a group that stops + # counting makes a resumed run apply a stale bias correction. + param = torch.nn.Parameter(torch.rand(4, dtype=torch.float, device="cuda")) + tst_optim = self.fused_optim( + [{"params": [param]}, {"params": []}], capturable=capturable, **self.options + ) + + num_steps = 3 + for _ in range(num_steps): + param.grad = torch.rand_like(param) + tst_optim.step() + + populated_step, empty_step = (int(g["step"]) for g in tst_optim.param_groups) + assert populated_step == num_steps + assert empty_step == populated_step + + # The counter is checkpointed through the param groups, so the empty group + # has to round trip with the same value as the populated one. + checkpoint = tst_optim.state_dict() + assert int(checkpoint["param_groups"][1]["step"]) == num_steps + def test_empty_param_at_end_of_group(self): tensors = [ torch.ones(4, dtype=torch.float, device="cuda"), diff --git a/tests/pytorch/test_fused_rope.py b/tests/pytorch/test_fused_rope.py index 50624df9e0..5cc4fc0f8a 100644 --- a/tests/pytorch/test_fused_rope.py +++ b/tests/pytorch/test_fused_rope.py @@ -143,7 +143,6 @@ def test_fused_rope_thd( start_positions: bool, margin: int, ) -> None: - device = torch.device("cuda:0") batch_size, head_num = 2, 64 cu_seqlens = [0, 400, 542, 711, 727, 752, 1270, 1426, 1450, 1954, 2044, 2048] @@ -179,7 +178,9 @@ def test_fused_rope_thd( t.requires_grad = True rotary_pos_emb = RotaryPositionEmbedding(hidden_size, rotary_percent, interleaved=interleaved) - emb = rotary_pos_emb(cu_seqlens_padded[-1]) + # A long frequency table with many packed spans exercises the default + # linear-grid dispatch heuristic without a test-only override. + emb = rotary_pos_emb(8192) assert emb.is_contiguous() for cp_rank in range(cp_size): @@ -345,7 +346,8 @@ def test_unfused_rope_thd_vs_bshd( grad_unfused_bshd.reshape(*grad_unfused_thd.shape), grad_unfused_thd ) torch.testing.assert_close( - grad_unfused_sbhd.transpose(1, 0).reshape(*grad_unfused_thd.shape), grad_unfused_thd + grad_unfused_sbhd.transpose(1, 0).reshape(*grad_unfused_thd.shape), + grad_unfused_thd, ) assert output_unfused_thd.is_contiguous() diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 09909d62d2..7048104e0b 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -25,6 +25,7 @@ OUTPUT_BUFFER_KEY, GRAD_INPUT_BUFFER_KEY, ) +from transformer_engine.pytorch.ops.fuser import OperationFuser from transformer_engine.pytorch._extra_state import UNSAFE_PICKLE_EXTRA_STATE_ENV from transformer_engine.pytorch.ops.fused import ( @@ -38,6 +39,7 @@ ) from transformer_engine.pytorch import ( QuantizedTensor, + Float8BlockQuantizer, Float8CurrentScalingQuantizer, Float8Quantizer, MXFP8Quantizer, @@ -65,6 +67,9 @@ fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) +fp8_block_scaling_available, reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available( + return_reason=True +) # Supported data types _dtypes: list[torch.dtype] = [torch.float32, torch.float16] @@ -83,6 +88,8 @@ if nvfp4_available: _quantization_list.append("nvfp4") _quantization_list.append("nvfp4_4over6") +if fp8_block_scaling_available: + _quantization_list.append("fp8_block_scaling") @pytest.fixture(autouse=True, scope="function") @@ -117,6 +124,8 @@ def maybe_skip_quantization( and not nvfp4_available ): pytest.skip(reason_for_no_nvfp4) + if quantization == "fp8_block_scaling" and not fp8_block_scaling_available: + pytest.skip(reason_for_no_fp8_block_scaling) # Check dims if dims is not None: @@ -128,6 +137,9 @@ def maybe_skip_quantization( elif quantization == "mxfp8": if math.prod(dims[:-1]) % 32 != 0 or dims[-1] % 32 != 0: pytest.skip("MXFP8 GEMMs require dims that are divisible by 32") + elif quantization == "fp8_block_scaling": + if math.prod(dims[:-1]) % 128 != 0 or dims[-1] % 128 != 0: + pytest.skip("FP8 block scaling requires dims that are divisible by 128") elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht"): if math.prod(dims[:-1]) % 16 != 0 or dims[-1] % 16 != 0: pytest.skip("NVFP4 GEMMs require dims that are divisible by 16") @@ -193,6 +205,17 @@ def make_reference_and_test_tensors( test = quantizer(test) elif quantization == "mxfp8": test = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3)(test) + elif quantization == "fp8_block_scaling": + tensor_type = "input" + if quantizer_role is not None: + tensor_type = quantizer_role.tensor_type + # Weights use 2D (128x128) blocks; everything else 1D (1x128). + test = Float8BlockQuantizer( + fp8_dtype=te.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=2 if tensor_type == "weight" else 1, + )(test) elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_rht"): tensor_type = "input" if quantizer_role is not None: @@ -422,6 +445,540 @@ def test_extra_tensors(self, size: int = 16) -> None: torch.testing.assert_close(x4, x4_orig + x3) +class _DualExtraOutput(te_ops.BasicOperation): + """Test helper: one op with two scaled extra outputs.""" + + num_extra_outputs = 2 + + def __init__(self, scales: tuple[float, float] = (1.0, 1.0)) -> None: + super().__init__() + self._scales = scales + + def op_forward(self, *args, **kwargs): + raise RuntimeError("_DualExtraOutput uses fuser_forward") + + def op_backward(self, *args, **kwargs): + raise RuntimeError("_DualExtraOutput uses fuser_backward") + + def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused): + del basic_op_ctxs, basic_op_extra_inputs + s0, s1 = self._scales + return input_, [(s0 * input_, s1 * input_)] + + def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs): + del basic_op_ctxs + s0, s1 = self._scales + g0, g1 = basic_op_grad_extra_outputs[0] + grad_extra = torch.zeros_like(grad_output) + if g0 is not None: + grad_extra = grad_extra + s0 * g0 + if g1 is not None: + grad_extra = grad_extra + s1 * g1 + return grad_output + grad_extra, [()], [()] + + +class TestExtraTensorChannels: + """Error handling and grad coverage for named extra-tensor channels.""" + + @pytest.mark.parametrize("with_extra_grad", (True, False)) + def test_internal_residual_connection( + self, + with_extra_grad: bool, + size: int = 16, + ) -> None: + """A channel can keep a residual connection inside a Sequential.""" + residual = te_ops.MakeExtraOutput() + body = te_ops.Bias(size=size, device="cpu") + add_residual = te_ops.AddExtraInput() + residual.set_extra_output_channel(0, "residual") + add_residual.set_extra_input_channel(0, "residual") + + model = te_ops.Sequential(residual, body, add_residual) + x = torch.rand((size,), requires_grad=True) + y, residual_out = model(x) + + torch.testing.assert_close(y, 2 * x + body.bias) + torch.testing.assert_close(residual_out, x) + dy = torch.rand_like(y) + if with_extra_grad: + dresidual = torch.rand_like(residual_out) + torch.autograd.backward((y, residual_out), (dy, dresidual)) + expected_dx = 2 * dy + dresidual + else: + y.backward(dy) + expected_dx = 2 * dy + torch.testing.assert_close(x.grad, expected_dx) + torch.testing.assert_close(body.bias.grad, dy) + + @pytest.mark.parametrize("output_to_caller", (True, False)) + def test_unconsumed_extra_output( + self, + output_to_caller: bool, + size: int = 16, + ) -> None: + """Unused public or hidden-unconsumed MakeExtraOutput treats None grad as zero.""" + residual = te_ops.MakeExtraOutput() + body = te_ops.Bias(size=size, device="cpu") + if not output_to_caller: + residual.set_extra_output_channel(0, "unused", output_to_caller=False) + + model = te_ops.Sequential(residual, body) + x = torch.rand((size,), requires_grad=True) + outputs = model(x) + if output_to_caller: + y, residual_out = outputs + torch.testing.assert_close(residual_out, x) + else: + assert isinstance(outputs, torch.Tensor) + y = outputs + + torch.testing.assert_close(y, x + body.bias) + dy = torch.rand_like(y) + y.backward(dy) + torch.testing.assert_close(x.grad, dy) + torch.testing.assert_close(body.bias.grad, dy) + + @pytest.mark.parametrize("fusion_kind", ("forward", "backward", "forward_backward")) + @pytest.mark.parametrize("output_to_caller", (True, False)) + def test_fused_internal_residual_connection( + self, + fusion_kind: str, + output_to_caller: bool, + size: int = 16, + ) -> None: + """Forward, backward, and joint fusions can own an internal channel.""" + + class FusedResidual(te_ops.FusedOperation): + """Fuse MakeExtraOutput, Bias, and AddExtraInput.""" + + _enabled = True + + def __init__(self, residual, body, add_residual) -> None: + super().__init__((residual, body, add_residual)) + + def fuser_forward( + self, + basic_op_ctxs, + input_, + *, + basic_op_extra_inputs, + **unused, + ): + del basic_op_ctxs + # The consumer slot is internal to this fusion, so the + # OperationFuser deliberately leaves it unset. + assert basic_op_extra_inputs[2][0] is None + residual_out = input_ if output_to_caller else None + return 2 * input_ + self.basic_ops[1].bias, [(residual_out,), (), ()] + + def fuser_backward( + self, + basic_op_ctxs, + grad_output, + *, + basic_op_grad_extra_outputs, + ): + del basic_op_ctxs + # The fusion owns the internal residual edge. The fuser also + # supplies a gradient when the residual is a public output. + grad_residual = basic_op_grad_extra_outputs[0][0] + return ( + 2 * grad_output + + (torch.zeros_like(grad_output) if grad_residual is None else grad_residual), + [(), (grad_output,), ()], + [(), (), (grad_output,)], + ) + + def fuse_residual(ops, **unused): + if not FusedResidual._enabled: + return ops + if ( + len(ops) == 3 + and isinstance(ops[0], te_ops.MakeExtraOutput) + and isinstance(ops[1], te_ops.Bias) + and isinstance(ops[2], te_ops.AddExtraInput) + ): + # We want to enable this fusion just for this test. + # Hence disable it after fusing it once in the test. + FusedResidual._enabled = False + return [FusedResidual(*ops)] + return ops + + residual = te_ops.MakeExtraOutput() + body = te_ops.Bias(size=size, device="cpu") + add_residual = te_ops.AddExtraInput() + residual.set_extra_output_channel( + 0, + "residual", + output_to_caller=output_to_caller, + ) + add_residual.set_extra_input_channel(0, "residual") + model = te_ops.Sequential(residual, body, add_residual) + + if fusion_kind == "forward": + te_ops.register_forward_fusion(fuse_residual, prepend=True) + elif fusion_kind == "backward": + te_ops.register_backward_fusion(fuse_residual, prepend=True) + else: + te_ops.register_forward_backward_fusion(fuse_residual, prepend=True) + x = torch.rand((size,), requires_grad=True) + outputs = model(x) + if output_to_caller: + y, residual_out = outputs + else: + assert isinstance(outputs, torch.Tensor) + y = outputs + + forward_ops = model._module_groups[0]._forward_ops + backward_ops = model._module_groups[0]._backward_ops + if fusion_kind in ("forward", "forward_backward"): + assert len(forward_ops) == 1 + assert isinstance(forward_ops[0][0], FusedResidual) + else: + assert len(forward_ops) == 3 + if fusion_kind in ("backward", "forward_backward"): + assert len(backward_ops) == 1 + assert isinstance(backward_ops[0][0], FusedResidual) + else: + assert len(backward_ops) == 3 + if fusion_kind == "forward_backward": + assert backward_ops[0][0] is forward_ops[0][0] + torch.testing.assert_close(y, 2 * x + body.bias) + dy = torch.rand_like(y) + if output_to_caller: + dresidual = torch.rand_like(residual_out) + torch.autograd.backward((y, residual_out), (dy, dresidual)) + expected_dx = 2 * dy + dresidual + else: + y.backward(dy) + expected_dx = 2 * dy + torch.testing.assert_close(x.grad, expected_dx) + torch.testing.assert_close(body.bias.grad, dy) + + def test_internal_extra_tensor_channel_fanout(self, size: int = 16) -> None: + """An internal extra output can feed multiple later consumers.""" + producer = te_ops.MakeExtraOutput() + consumer1 = te_ops.AddExtraInput() + consumer2 = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + consumer1.set_extra_input_channel(0, "route") + consumer2.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, consumer1, consumer2) + + x = torch.rand((size,), requires_grad=True) + y, route = model(x) + + # Main path: x -> x + route -> x + route + route. + torch.testing.assert_close(y, 3 * x) + torch.testing.assert_close(route, x) + dy = torch.rand_like(y) + droute = torch.rand_like(route) + torch.autograd.backward((y, route), (dy, droute)) + # The channel fan-out contributes two independent gradient paths. + torch.testing.assert_close(x.grad, 3 * dy + droute) + + # Internal slots are unavailable before forward, so grad discovery + # must tolerate them when no public input requires gradients. + x_no_grad = x.detach() + y_no_grad, route_no_grad = model(x_no_grad) + torch.testing.assert_close(y_no_grad, 3 * x_no_grad) + torch.testing.assert_close(route_no_grad, x_no_grad) + + def test_internal_extra_tensor_channel_can_be_hidden(self, size: int = 16) -> None: + """A non-public channel still propagates forward and backward.""" + producer = te_ops.MakeExtraOutput() + consumer1 = te_ops.AddExtraInput() + consumer2 = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route", output_to_caller=False) + consumer1.set_extra_input_channel(0, "route") + consumer2.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, consumer1, consumer2) + + x = torch.rand((size,), requires_grad=True) + y = model(x) + assert isinstance(y, torch.Tensor) + torch.testing.assert_close(y, 3 * x) + + dy = torch.rand_like(y) + y.backward(dy) + torch.testing.assert_close(x.grad, 3 * dy) + + def test_internal_and_external_extra_tensor_inputs(self, size: int = 16) -> None: + """Unbound slots remain public when other slots use internal channels.""" + producer = te_ops.MakeExtraOutput() + internal_consumer = te_ops.AddExtraInput() + external_consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + internal_consumer.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, internal_consumer, external_consumer) + + x = torch.rand((size,), requires_grad=True) + extra = torch.rand((size,), requires_grad=True) + y, route = model(x, extra) + + torch.testing.assert_close(y, 2 * x + extra) + torch.testing.assert_close(route, x) + dy = torch.rand_like(y) + y.backward(dy) + torch.testing.assert_close(x.grad, 2 * dy) + torch.testing.assert_close(extra.grad, dy) + + def test_named_extra_input_requires_producer(self) -> None: + """A named input cannot fall back to a caller-provided tensor.""" + consumer = te_ops.AddExtraInput() + consumer.set_extra_input_channel(0, "missing") + with pytest.raises(ValueError, match="has no producer"): + OperationFuser([consumer]) + + def test_consumer_before_producer(self) -> None: + """Channels only connect forward; a later producer does not satisfy an earlier consumer.""" + consumer = te_ops.AddExtraInput() + producer = te_ops.MakeExtraOutput() + consumer.set_extra_input_channel(0, "route") + producer.set_extra_output_channel(0, "route") + with pytest.raises(ValueError, match="has no earlier producer"): + OperationFuser([consumer, producer]) + + def test_set_extra_channel_rejects_invalid_index(self) -> None: + """Slot indices must be in range; negatives and OOB are rejected at bind time.""" + producer = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + + with pytest.raises(IndexError, match="out of range"): + producer.set_extra_output_channel(-1, "route") + with pytest.raises(IndexError, match="out of range"): + producer.set_extra_output_channel(1, "route") + with pytest.raises(IndexError, match="out of range"): + consumer.set_extra_input_channel(-1, "route") + with pytest.raises(IndexError, match="out of range"): + consumer.set_extra_input_channel(1, "route") + + def test_set_extra_channel_rejects_invalid_name(self) -> None: + """Channel names must be non-empty strings.""" + producer = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + + with pytest.raises(ValueError, match="non-empty string"): + producer.set_extra_output_channel(0, "") + with pytest.raises(ValueError, match="non-empty string"): + consumer.set_extra_input_channel(0, "") + with pytest.raises(ValueError, match="non-empty string"): + producer.set_extra_output_channel(0, 123) # type: ignore[arg-type] + with pytest.raises(TypeError, match="output_to_caller must be a bool"): + producer.set_extra_output_channel( + 0, + "route", + output_to_caller=1, # type: ignore[arg-type] + ) + + def test_extra_channels_lock_after_capture(self, size: int = 16) -> None: + """Channel bindings freeze once a fuser captures them.""" + producer = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + consumer.set_extra_input_channel(0, "route") + assert not producer._extra_tensor_channels_locked + + # Construction alone locks channels, before any forward pass. + OperationFuser([producer, consumer]) + assert producer._extra_tensor_channels_locked + assert consumer._extra_tensor_channels_locked + with pytest.raises(RuntimeError, match="already captured its channel routing"): + producer.set_extra_output_channel(0, None) + + # Fresh ops for the Sequential path. + producer = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + consumer.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, consumer) + + x = torch.rand((size,)) + y, route = model(x) + torch.testing.assert_close(y, 2 * x) + torch.testing.assert_close(route, x) + + with pytest.raises(RuntimeError, match="already captured its channel routing"): + consumer.set_extra_input_channel(0, None) + with pytest.raises(RuntimeError, match="already captured its channel routing"): + producer.set_extra_output_channel(0, "route", output_to_caller=False) + + @pytest.mark.parametrize("layout", ("two_ops", "same_op")) + def test_duplicate_extra_output_channel_names(self, layout: str) -> None: + """A channel name may have at most one producer, across ops or slots.""" + consumer = te_ops.AddExtraInput() + consumer.set_extra_input_channel(0, "route") + if layout == "two_ops": + producer1 = te_ops.MakeExtraOutput() + producer2 = te_ops.MakeExtraOutput() + producer1.set_extra_output_channel(0, "route") + producer2.set_extra_output_channel(0, "route") + ops = [producer1, producer2, consumer] + else: + producer = _DualExtraOutput() + producer.set_extra_output_channel(0, "route") + producer.set_extra_output_channel(1, "route") + ops = [producer, consumer] + with pytest.raises(ValueError, match="multiple producers"): + OperationFuser(ops) + + def test_one_extra_input_has_single_source(self, size: int = 16) -> None: + """Rebinding selects one source and leaves the other output public.""" + producer_a = te_ops.MakeExtraOutput() + producer_b = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + producer_a.set_extra_output_channel(0, "a") + producer_b.set_extra_output_channel(0, "b") + consumer.set_extra_input_channel(0, "a") + consumer.set_extra_input_channel(0, "b") + fuser = OperationFuser([producer_a, producer_b, consumer]) + assert fuser._basic_op_extra_input_sources[2] == [(1, 0)] + assert fuser.num_extra_inputs == 0 + + x = torch.rand((size,)) + y, output_a, output_b = fuser(x) + torch.testing.assert_close(y, 2 * x) + torch.testing.assert_close(output_a, x) + torch.testing.assert_close(output_b, x) + + def test_mixed_public_and_hidden_channel_outputs(self, size: int = 16) -> None: + """Only configured public outputs are returned, in slot order.""" + producer = _DualExtraOutput(scales=(2.0, 3.0)) + consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "internal", output_to_caller=False) + producer.set_extra_output_channel(1, "public") + consumer.set_extra_input_channel(0, "internal") + model = te_ops.Sequential(producer, consumer) + + x = torch.rand((size,), requires_grad=True) + y, public = model(x) + torch.testing.assert_close(y, 3 * x) + torch.testing.assert_close(public, 3 * x) + + dy = torch.rand_like(y) + dpublic = torch.rand_like(public) + torch.autograd.backward((y, public), (dy, dpublic)) + torch.testing.assert_close(x.grad, 3 * dy + 3 * dpublic) + + @pytest.mark.parametrize("output_to_caller", (True, False)) + def test_fused_op_cannot_omit_required_channel_output( + self, + output_to_caller: bool, + size: int = 16, + ) -> None: + """A fusion must materialize public outputs and cross-fusion channels.""" + + class FusedProducer(te_ops.FusedOperation): + _enabled = True + + def __init__(self, producer) -> None: + super().__init__((producer,)) + + def fuser_forward(self, basic_op_ctxs, input_, **unused): + del basic_op_ctxs + return input_, [(None,)] + + def fuse_producer(ops, **unused): + if ( + FusedProducer._enabled + and len(ops) == 2 + and isinstance(ops[0], te_ops.MakeExtraOutput) + and isinstance(ops[1], te_ops.AddExtraInput) + ): + FusedProducer._enabled = False + return [FusedProducer(ops[0]), ops[1]] + return ops + + producer = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel( + 0, + "route", + output_to_caller=output_to_caller, + ) + consumer.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, consumer) + te_ops.register_forward_fusion(fuse_producer, prepend=True) + + x = torch.rand((size,), requires_grad=True) + error = "is public" if output_to_caller else "outside its forward fusion" + with pytest.raises(RuntimeError, match=error): + model(x) + + def test_fresh_internal_output_preserves_grad_requirement(self) -> None: + """A freshly computed internal channel still receives a consumer gradient.""" + + # A BasicOperation with one extra output that is freshly computed instead of + # retrieved from a previous op's tensor. + class MakeScale(te_ops.BasicOperation): + num_extra_outputs = 1 + + def op_forward(self, *args, **kwargs): + raise RuntimeError("MakeScale uses fuser_forward") + + def op_backward(self, *args, **kwargs): + raise RuntimeError("MakeScale uses fuser_backward") + + def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused): + del basic_op_extra_inputs + basic_op_ctxs[0].save_for_backward(input_) + return input_, [(input_.square().mean(dim=-1),)] + + def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs): + (input_,) = basic_op_ctxs[0].saved_tensors + grad_scale = basic_op_grad_extra_outputs[0][0] + assert grad_scale is not None + grad_input = grad_output + grad_scale.unsqueeze(-1) * 2 * input_ / input_.size(-1) + return grad_input, [()], [()] + + class ScaleByExtra(te_ops.BasicOperation): + num_extra_inputs = 1 + + def op_forward(self, *args, **kwargs): + raise RuntimeError("ScaleByExtra uses fuser_forward") + + def op_backward(self, *args, **kwargs): + raise RuntimeError("ScaleByExtra uses fuser_backward") + + def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused): + scale = basic_op_extra_inputs[0][0] + ctx = basic_op_ctxs[0] + # Match scaled activations: only compute scale grads when the + # fuser marked this fresh internal channel as requiring grad. + ctx.extra_input_requires_grad = scale.requires_grad + ctx.save_for_backward(input_, scale) + return input_ * scale.unsqueeze(-1), [()] + + def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs): + del basic_op_grad_extra_outputs + ctx = basic_op_ctxs[0] + input_, scale = ctx.saved_tensors + grad_input = grad_output * scale.unsqueeze(-1) + grad_scale = ( + (grad_output * input_).sum(dim=-1) if ctx.extra_input_requires_grad else None + ) + return grad_input, [()], [(grad_scale,)] + + producer = MakeScale() + consumer = ScaleByExtra() + producer.set_extra_output_channel(0, "scale", output_to_caller=False) + consumer.set_extra_input_channel(0, "scale") + model = te_ops.Sequential(producer, consumer) + + x_ref = torch.randn((5, 8), requires_grad=True) + x_test = x_ref.detach().clone().requires_grad_(True) + scale_ref = x_ref.square().mean(dim=-1) + y_ref = x_ref * scale_ref.unsqueeze(-1) + y_test = model(x_test) + assert isinstance(y_test, torch.Tensor) + torch.testing.assert_close(y_test, y_ref) + + dy = torch.rand_like(y_ref) + y_ref.backward(dy) + y_test.backward(dy) + torch.testing.assert_close(x_test.grad, x_ref.grad) + + class TestFuser: """Tests for operation fusion infrastructure""" @@ -1036,8 +1593,9 @@ def _test_basic_linear( ) torch.testing.assert_close(dw_test, w_ref.grad, **tols) - @pytest.mark.parametrize("weight_shape", ((64, 32), (3, 5))) - @pytest.mark.parametrize("in_shape", ((-1,), (5, 1, -1), (4, 2, 4, -1))) + # (128, 128) + a 128-token in_shape keep FP8 block scaling (128-divisible dims) unskipped. + @pytest.mark.parametrize("weight_shape", ((64, 32), (3, 5), (128, 128))) + @pytest.mark.parametrize("in_shape", ((-1,), (5, 1, -1), (4, 2, 4, -1), (128, -1))) @pytest.mark.parametrize("dtype", _dtypes) @pytest.mark.parametrize("quantization", _quantization_list) @pytest.mark.parametrize("accumulate_into_main_grad", (False, True)) @@ -2972,6 +3530,7 @@ def test_backward_activation_bias( @pytest.mark.parametrize("in_shape", ((-1,), (6, 16, -1))) @pytest.mark.parametrize("dtype", _dtypes) @pytest.mark.parametrize("zero_centered_gamma", (False, True)) + @pytest.mark.parametrize("with_extra_grad", (True, False)) def test_backward_add_rmsnorm( self, *, @@ -2981,6 +3540,7 @@ def test_backward_add_rmsnorm( device: torch.device = "cuda", eps: float = 0.3, zero_centered_gamma: bool, + with_extra_grad: bool, ) -> None: """Fused backward RMNorm + add""" @@ -3019,7 +3579,10 @@ def test_backward_add_rmsnorm( else: y1_ref = x_ref / torch.sqrt(eps + var_ref) * w_ref y2_ref = x_ref - (y1_ref * dy1_ref + y2_ref * dy2_ref).sum().backward() + if with_extra_grad: + (y1_ref * dy1_ref + y2_ref * dy2_ref).sum().backward() + else: + (y1_ref * dy1_ref).sum().backward() # Implementation with fusible operations model = te_ops.Sequential( @@ -3036,7 +3599,10 @@ def test_backward_add_rmsnorm( model[1].weight.copy_(w_test) del w_test y1_test, y2_test = model(x_test) - (y1_test * dy1_test + y2_test * dy2_test).sum().backward() + if with_extra_grad: + (y1_test * dy1_test + y2_test * dy2_test).sum().backward() + else: + (y1_test * dy1_test).sum().backward() # Check that backward operations have been fused backward_ops = model._module_groups[0]._backward_ops @@ -3058,6 +3624,7 @@ def test_backward_add_rmsnorm( @pytest.mark.parametrize("dtype", _dtypes) @pytest.mark.parametrize("quantization", _quantization_list) + @pytest.mark.parametrize("with_extra_grad", (True, False)) def test_backward_linear_add( self, *, @@ -3067,6 +3634,7 @@ def test_backward_linear_add( device: torch.device = "cuda", quantization: Optional[str], quantized_weight: bool = False, + with_extra_grad: bool, ) -> None: """Backward dgrad GEMM + add""" @@ -3114,7 +3682,10 @@ def test_backward_linear_add( # Plain PyTorch implementation y1_ref = torch.nn.functional.linear(x_ref, w_ref) y2_ref = x_ref - (y1_ref * dy1_ref + y2_ref * dy2_ref).sum().backward() + if with_extra_grad: + (y1_ref * dy1_ref + y2_ref * dy2_ref).sum().backward() + else: + (y1_ref * dy1_ref).sum().backward() # Implementation with fusible operations recipe = make_recipe(quantization) @@ -3134,7 +3705,10 @@ def test_backward_linear_add( del w_test with te.autocast(enabled=quantized_compute, recipe=recipe): y1_test, y2_test = model(x_test) - (y1_test * dy1_test + y2_test * dy2_test).sum().backward() + if with_extra_grad: + (y1_test * dy1_test + y2_test * dy2_test).sum().backward() + else: + (y1_test * dy1_test).sum().backward() # Check that backward operations have been fused backward_ops = model._module_groups[0]._backward_ops @@ -3394,7 +3968,7 @@ def test_layernorm_mlp( quantization: Optional[str], device: torch.device = "cuda", hidden_size: int = 256, - sequence_length: int = 48, + sequence_length: int = 64, batch_size: int = 4, ffn_hidden_size: int = 384, layernorm_epsilon: float = 1e-5, diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 8e630810de..048240df9e 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -17,6 +17,8 @@ import transformer_engine.pytorch as te from transformer_engine.common import recipe from transformer_engine.pytorch import ( + Float8BlockQuantizer, + Float8CurrentScalingQuantizer, Float8Quantizer, Fp8Padding, Fp8Unpadding, @@ -34,6 +36,11 @@ general_grouped_gemm, general_grouped_gemm_for_grouped_tensor, ) +from transformer_engine.pytorch.constants import TE_DType +from transformer_engine.pytorch.module.grouped_linear import ( + _GroupedLinear, + is_module_grouped_tensor_path_supported, +) from transformer_engine.pytorch.quantization import ( FP8GlobalStateManager, get_align_size_for_quantization, @@ -102,6 +109,23 @@ def nvfp4_row_scaled(): return nvfp4_recipe +def nvfp4_row_scaled_quantized_backward(): + # Same row-scaled activation recipe as nvfp4_row_scaled(), but with + # backward_override=None so the backward runs in NVFP4 instead of falling back + # to high precision. + nvfp4_recipe = recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + disable_2d_quantization=True, + row_scaled_activation=True, + backward_override=None, + ) + nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() + nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams() + nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams() + return nvfp4_recipe + + def nvfp4_4over6(): nvfp4_recipe = recipe.NVFP4BlockScaling( disable_rht=True, @@ -402,6 +426,84 @@ def test_grouped_linear_accuracy( torch.testing.assert_close(o, o_ref, rtol=rtol, atol=atol) +@pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4) +@pytest.mark.parametrize("dtype", [torch.bfloat16], ids=str) +@pytest.mark.parametrize("num_gemms", [1, 3]) +@pytest.mark.parametrize("bs", [2]) +@pytest.mark.parametrize("bias", all_boolean) +def test_grouped_linear_row_scaled_quantized_backward(dtype, num_gemms, bs, bias, model="126m"): + """Row-scaled NVFP4 GroupedLinear with quantized (non-fallback) NVFP4 backward. + + With ``backward_override=None`` the wgrad is computed in NVFP4: the row-scaled + activation becomes operand A of the ``NT`` grouped GEMM, which this PR routes + through the per-expert dense ``general_gemm`` loop. GroupedLinear must then + match a stack of independent dense ``Linear`` layers bit-for-bit, since both + execute the exact same per-expert quantize + GEMM kernels. + """ + recipe_row_scaled = nvfp4_row_scaled_quantized_backward() + config = model_configs[model] + if dtype not in get_nvfp4_inp_supported_dtypes(recipe_row_scaled, dtype): + pytest.skip(f"Input dtype {dtype} not supported for row-scaled NVFP4.") + + grouped_linear = ( + GroupedLinear( + num_gemms, + config.hidden_size, + 4 * config.hidden_size, + bias=bias, + params_dtype=dtype, + device="cuda", + ) + .cuda() + .eval() + ) + sequential_linear = torch.nn.ModuleList( + [ + Linear( + config.hidden_size, + 4 * config.hidden_size, + bias=bias, + params_dtype=dtype, + device="cuda", + ).eval() + for _ in range(num_gemms) + ] + ) + + # Share weights/biases so the two paths are numerically comparable. + with torch.no_grad(): + for i in range(num_gemms): + sequential_linear[i].weight = Parameter(getattr(grouped_linear, f"weight{i}").clone()) + if bias: + sequential_linear[i].bias = Parameter(getattr(grouped_linear, f"bias{i}").clone()) + + outputs_ref = _test_grouped_linear_accuracy( + sequential_linear, + num_gemms, + bs, + dtype, + config, + recipe_row_scaled, + fp8=True, + fuse_wgrad_accumulation=False, + ) + outputs = _test_grouped_linear_accuracy( + grouped_linear, + num_gemms, + bs, + dtype, + config, + recipe_row_scaled, + fp8=True, + fuse_wgrad_accumulation=False, + ) + + # GroupedLinear is a per-expert loop over the same dense kernels, so the + # forward output, dgrad, and (row-scaled) wgrad must match bit-for-bit. + for o, o_ref in zip(outputs, outputs_ref): + torch.testing.assert_close(o, o_ref, rtol=0, atol=0) + + @pytest.mark.skipif( torch.cuda.get_device_capability() != (9, 0), reason="Only enable CUTLASS grouped gemm on Hopper", @@ -1405,46 +1507,124 @@ def test_grouped_gemm_grouped_tensor(z, m, n, k, case, layout, accumulate, use_b torch.testing.assert_close(o, o_ref, **tols) +@pytest.mark.parametrize("use_bias_scale", [False, True]) +def test_grouped_gemm_grouped_tensor_zero_work_bias(use_bias_scale) -> None: + """A grouped bias operation is a no-op when every group has zero rows. + + Zero-sized CUDA tensors may legally have a null data pointer. Exercise both bias entry points + so neither the ordinary nor scaled path mistakes that pointer for missing output storage. + This BF16 case runs on both Hopper and Blackwell when grouped cuBLASLt GEMM is available. + """ + if not is_module_grouped_tensor_path_supported( + None, + torch.bfloat16, + ): + pytest.skip("BF16 GroupedTensor GEMM is unavailable.") + + num_groups = 4 + in_features = 256 + out_features = 256 + m_sizes = [0] * num_groups + dtype = torch.bfloat16 + device = torch.device("cuda") + + weights = [ + torch.randn(out_features, in_features, dtype=dtype, device=device) + for _ in range(num_groups) + ] + biases = [torch.randn(1, out_features, dtype=dtype, device=device) for _ in range(num_groups)] + grouped_weights = _make_grouped_tensor_uniform( + num_groups, out_features, in_features, device, dtype + ) + grouped_input = _make_grouped_tensor_from_splits(m_sizes, in_features, device, dtype) + grouped_output = _make_grouped_tensor_from_splits(m_sizes, out_features, device, dtype) + grouped_bias = _make_grouped_tensor_uniform(num_groups, 1, out_features, device, dtype) + _pack_grouped_tensor(grouped_weights, weights) + _pack_grouped_tensor(grouped_bias, biases) + + bias_scale = torch.empty(0, dtype=torch.float32, device=device) if use_bias_scale else None + general_grouped_gemm_for_grouped_tensor( + grouped_weights, + grouped_input, + grouped_output, + layout="TN", + bias=grouped_bias, + bias_scale=bias_scale, + ) + torch.cuda.synchronize() + + assert grouped_output.rowwise_data.numel() == 0 + + @pytest.mark.parametrize("layout", ["TN", "NN", "NT"]) @pytest.mark.parametrize("accumulate", [False, True]) -@pytest.mark.parametrize("quant_type", ["bf16", "mxfp8"]) +@pytest.mark.parametrize( + "quant_type", ["bf16", "fp8_current_scaling", "mxfp8", "fp8_block_scaling"] +) def test_grouped_gemm_grouped_tensor_zero_work(layout, accumulate, quant_type) -> None: """Grouped GEMM with all-zero split sizes (zero total work). For wgrad (NT layout) the output should be zero when not accumulating, or unchanged when accumulating with beta=1. """ - if torch.cuda.get_device_capability() < (10, 0): - pytest.skip("Grouped GEMM requires Blackwell (SM100) or newer.") if not is_bf16_available(): pytest.skip("bfloat16 is required for grouped GEMM test.") - if quant_type == "mxfp8" and not mxfp8_available: - pytest.skip(reason_for_no_mxfp8) z = 4 k, n = 256, 256 dtype = torch.bfloat16 device = torch.device("cuda") - use_mxfp8 = quant_type == "mxfp8" + + test_recipe = { + "bf16": None, + "fp8_current_scaling": recipe.Float8CurrentScaling(), + "mxfp8": recipe.MXFP8BlockScaling(), + "fp8_block_scaling": recipe.Float8BlockScaling(), + }[quant_type] + if not is_module_grouped_tensor_path_supported( + test_recipe, + dtype, + ): + pytest.skip(f"{quant_type} grouped-tensor GEMM is unavailable") transa = layout[0] == "T" transb = layout[1] == "T" zero_first_dims = torch.zeros(z, dtype=torch.int64, device=device) + def _make_quantizer(fp8_dtype, rowwise, columnwise): + if quant_type == "fp8_current_scaling": + quantizer = Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device=device) + quantizer.set_usage(rowwise=rowwise, columnwise=columnwise) + elif quant_type == "mxfp8": + quantizer = MXFP8Quantizer( + fp8_dtype=fp8_dtype, + rowwise=rowwise, + columnwise=columnwise, + ) + elif quant_type == "fp8_block_scaling": + quantizer = Float8BlockQuantizer( + fp8_dtype=fp8_dtype, + rowwise=rowwise, + columnwise=columnwise, + force_pow_2_scales=False, + amax_epsilon=0.0, + block_scaling_dim=1, + ) + else: + raise ValueError(f"Unsupported quantized zero-work test type {quant_type}") + quantizer.optimize_for_gemm = True + return quantizer + def _make_zero_tokens_grouped_tensor(logical_last_dim, is_a): """Create a GroupedTensor with non-zero logical_shape but zero first_dims.""" buf = torch.randn(0, logical_last_dim, dtype=dtype, device=device) - if use_mxfp8: + if test_recipe is not None: if is_a: rowwise, columnwise = transa, not transa else: rowwise, columnwise = not transb, transb - quantizer = MXFP8Quantizer( - fp8_dtype=tex.DType.kFloat8E4M3, - rowwise=rowwise, - columnwise=columnwise, - ) - quantizer.optimize_for_gemm = True + fp8_dtype = TE_DType[torch.float8_e4m3fn] + quantizer = _make_quantizer(fp8_dtype, rowwise, columnwise) return tex.group_quantize(buf, quantizer, z, zero_first_dims) return GroupedTensor.make_grouped_tensor( num_tensors=z, @@ -1459,12 +1639,18 @@ def _make_zero_tokens_grouped_tensor(logical_last_dim, is_a): if layout in ("TN", "NN"): weight_tensors = [torch.randn(n, k, dtype=dtype, device=device) for _ in range(z)] - if use_mxfp8: - grouped_A = _make_grouped_tensor_quantized_mxfp8( - weight_tensors, + if test_recipe is not None: + grouped_weight = torch.cat(weight_tensors, dim=0) + weight_quantizer = _make_quantizer( + TE_DType[torch.float8_e4m3fn], rowwise=transa, columnwise=not transa, - device=device, + ) + grouped_A = tex.group_quantize( + grouped_weight, + weight_quantizer, + z, + torch.full((z,), n, dtype=torch.int64, device=device), ) else: grouped_A = _make_grouped_tensor_uniform(z, n, k, device, dtype) @@ -1752,20 +1938,645 @@ def test_fp8_grouped_gemm(shape, accumulate): _fp8_available, _reason_for_no_fp8 = fp8_available, reason_for_no_fp8 _mxfp8_available, _reason_for_no_mxfp8 = mxfp8_available, reason_for_no_mxfp8 _nvfp4_available, _reason_for_no_nvfp4 = nvfp4_available, reason_for_no_nvfp4 +_fp8_block_scaling_available, _reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available( + return_reason=True +) @pytest.fixture(autouse=True) def _reset_fp8_state(monkeypatch): - monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "0") + monkeypatch.delenv(_FUSED_GROUPED_GEMM_ENV, raising=False) yield FP8GlobalStateManager.reset() monkeypatch.delenv(_FUSED_GROUPED_GEMM_ENV, raising=False) +@pytest.mark.parametrize( + "m_splits,exception", + [([256, 256], ValueError), (torch.tensor([256, 256]), ValueError)], + ids=["python-list", "cpu-tensor"], +) +def test_single_grouped_weight_rejects_host_m_splits(monkeypatch, m_splits, exception): + """A single parent parameter must never fall back to host-split per-expert GEMMs.""" + if not is_module_grouped_tensor_path_supported( + None, + torch.bfloat16, + ): + pytest.skip("Native GroupedTensor GEMM is unavailable on this system.") + + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + grouped_linear = GroupedLinear( + 2, + 64, + 64, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + single_grouped_weight=True, + use_grouped_tensor=True, + ) + x = torch.randn(512, 64, dtype=torch.bfloat16, device="cuda") + with pytest.raises(exception, match="requires.*CUDA"): + grouped_linear(x, m_splits) + + +@pytest.mark.parametrize( + "fp8_recipe", + [ + pytest.param(None, id="bf16"), + pytest.param( + recipe.Float8CurrentScaling(), + marks=pytest.mark.skipif(not _fp8_available, reason=_reason_for_no_fp8), + id="fp8-current-scaling", + ), + pytest.param( + recipe.MXFP8BlockScaling(), + marks=pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8), + id="mxfp8", + ), + pytest.param( + recipe.Float8BlockScaling(), + marks=pytest.mark.skipif( + not _fp8_block_scaling_available, + reason=_reason_for_no_fp8_block_scaling, + ), + id="fp8-block-scaling", + ), + ], +) +def test_single_grouped_weight_matches_discrete_grouped_tensor_path(monkeypatch, fp8_recipe): + """Match single and discrete weights while both use CUDA m_splits and grouped GEMM.""" + if not is_module_grouped_tensor_path_supported( + fp8_recipe, + torch.bfloat16, + ): + pytest.skip("Recipe is not supported with a single grouped weight on this system.") + + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + FP8GlobalStateManager.reset() + + num_gemms = 3 + in_features = 256 + out_features = 256 + m_splits = torch.tensor([256, 512, 256], dtype=torch.int64, device="cuda") + total_tokens = int(m_splits.sum()) + weights = torch.randn( + num_gemms, + out_features, + in_features, + dtype=torch.bfloat16, + device="cuda", + ) + + discrete = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + use_grouped_tensor=True, + ) + single = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + single_grouped_weight=True, + use_grouped_tensor=True, + ) + with torch.no_grad(): + for idx in range(num_gemms): + getattr(discrete, f"weight{idx}").copy_(weights[idx]) + single.weight.rowwise_data.view_as(weights).copy_(weights) + + x = torch.randn(total_tokens, in_features, dtype=torch.bfloat16, device="cuda") + dy = torch.randn(total_tokens, out_features, dtype=torch.bfloat16, device="cuda") + x_discrete = x.detach().clone().requires_grad_(True) + x_single = x.detach().clone().requires_grad_(True) + + with autocast(enabled=fp8_recipe is not None, recipe=fp8_recipe): + y_discrete = discrete(x_discrete, m_splits) + y_single = single(x_single, m_splits) + y_discrete.backward(dy) + y_single.backward(dy) + + tolerances = dict(rtol=1e-2, atol=5e-3) + torch.testing.assert_close(y_single.float(), y_discrete.float(), **tolerances) + torch.testing.assert_close(x_single.grad.float(), x_discrete.grad.float(), **tolerances) + discrete_wgrad = torch.stack( + [getattr(discrete, f"weight{idx}").grad for idx in range(num_gemms)] + ) + torch.testing.assert_close(single.weight.grad.float(), discrete_wgrad.float(), **tolerances) + + +@pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8) +def test_single_grouped_weight_mxfp8_workspace_cache(monkeypatch): + """BF16 primary weights update one persistent MXFP8 grouped workspace per iteration.""" + mxfp8_recipe = recipe.MXFP8BlockScaling() + if not is_module_grouped_tensor_path_supported( + mxfp8_recipe, + torch.bfloat16, + ): + pytest.skip("MXFP8 single-weight GroupedTensor path is unavailable on this system.") + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + FP8GlobalStateManager.reset() + grouped_linear = GroupedLinear( + 2, + 256, + 256, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + single_grouped_weight=True, + use_grouped_tensor=True, + ) + + x = torch.randn(512, 256, dtype=torch.bfloat16, device="cuda", requires_grad=True) + m_splits = torch.tensor([256, 256], dtype=torch.int64, device="cuda") + + with autocast(enabled=True, recipe=mxfp8_recipe): + grouped_linear(x, m_splits, is_first_microbatch=True) + workspace = grouped_linear._fp8_workspaces["weight"] + assert isinstance(workspace, GroupedTensor) + pointers = ( + workspace.rowwise_data.data_ptr(), + workspace.columnwise_data.data_ptr(), + workspace.scale_inv.data_ptr(), + workspace.columnwise_scale_inv.data_ptr(), + ) + old_data = workspace.rowwise_data.clone() + + with torch.no_grad(): + grouped_linear.weight.rowwise_data.add_(1) + grouped_linear(x, m_splits, is_first_microbatch=False) + assert torch.equal(workspace.rowwise_data, old_data) + + grouped_linear(x, m_splits, is_first_microbatch=True) + assert pointers == ( + workspace.rowwise_data.data_ptr(), + workspace.columnwise_data.data_ptr(), + workspace.scale_inv.data_ptr(), + workspace.columnwise_scale_inv.data_ptr(), + ) + assert not torch.equal(workspace.rowwise_data, old_data) + + +@pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8) +@pytest.mark.parametrize("fp8_recipe", [recipe.MXFP8BlockScaling()], ids=recipe_id) +def test_single_grouped_weight_with_disabled_weight_preswizzle(monkeypatch, fp8_recipe): + """Grouped weight preparation preserves a disabled preswizzle decision.""" + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + FP8GlobalStateManager.reset() + with quantized_model_init(enabled=True, recipe=fp8_recipe): + grouped_linear = GroupedLinear( + 2, + 256, + 256, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + single_grouped_weight=True, + use_grouped_tensor=True, + ) + + weight_quantizer = grouped_linear.weight.quantizer + assert weight_quantizer is not None + preswizzle = grouped_linear._enable_weight_preswizzle( + weight_quantizer, + grouped_linear.weight, + ) + assert preswizzle is False + weight_quantizer.optimize_for_gemm = preswizzle + grouped_weight, new_workspaces = _GroupedLinear._prepare_weights_for_grouped_tensor_gemm( + (grouped_linear.weight,), + [weight_quantizer], + [None], + num_gemms=grouped_linear.num_gemms, + single_grouped_weight=True, + with_quantized_compute=True, + columnwise_usage=True, + activation_dtype=torch.bfloat16, + is_first_microbatch=True, + skip_fp8_weight_update=None, + cache_weight=True, + ) + + assert weight_quantizer.optimize_for_gemm is False + assert len(new_workspaces) == 1 + assert new_workspaces[0] is None + assert grouped_weight is grouped_linear.weight + assert hasattr(grouped_weight, "_with_gemm_swizzled_scales") + assert grouped_weight._with_gemm_swizzled_scales is False + + +@pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8) +def test_single_grouped_primary_mxfp8_bypasses_weight_workspace(monkeypatch): + """An MXFP8 primary grouped parameter is already GEMM-ready and is not requantized.""" + mxfp8_recipe = recipe.MXFP8BlockScaling() + if not is_module_grouped_tensor_path_supported( + mxfp8_recipe, + torch.bfloat16, + ): + pytest.skip("MXFP8 single-weight GroupedTensor path is unavailable on this system.") + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + FP8GlobalStateManager.reset() + with quantized_model_init(enabled=True, recipe=mxfp8_recipe): + grouped_linear = GroupedLinear( + 2, + 256, + 256, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + single_grouped_weight=True, + use_grouped_tensor=True, + ) + + x = torch.randn(512, 256, dtype=torch.bfloat16, device="cuda") + m_splits = torch.tensor([256, 256], dtype=torch.int64, device="cuda") + with torch.no_grad(), autocast(enabled=True, recipe=mxfp8_recipe): + grouped_linear(x, m_splits, is_first_microbatch=True) + assert grouped_linear.weight.quantizer is not None + assert "weight" not in grouped_linear._fp8_workspaces + + def _clone_outputs(outputs): return [None if out is None else out.detach().clone() for out in outputs] +def _grouped_linear_weight_params(module): + if module.single_grouped_weight: + return [module.weight] + return [getattr(module, f"weight{i}") for i in range(module.num_gemms)] + + +def _grouped_linear_bias_params(module): + if not module.use_bias: + return [] + if module.single_grouped_bias: + return [module.bias] + return [getattr(module, f"bias{i}") for i in range(module.num_gemms)] + + +def _run_grouped_parameter_layout( + *, + use_grouped_tensor, + fp8_recipe, + single_grouped_weight, + single_grouped_bias, + use_bias, + delay_wgrad_compute, + fuse_wgrad_accumulation, + x_base, + dy, + weights, + biases, + m_splits, + save_original_input=False, +): + """Run one layout and return all numerically observable forward/backward results.""" + FP8GlobalStateManager.reset() + num_gemms, out_features, in_features = weights.shape + module = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=use_bias, + params_dtype=torch.bfloat16, + device="cuda", + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + delay_wgrad_compute=delay_wgrad_compute, + single_grouped_weight=single_grouped_weight, + single_grouped_bias=single_grouped_bias, + use_grouped_tensor=use_grouped_tensor, + save_original_input=save_original_input, + ) + + with torch.no_grad(): + if single_grouped_weight: + module.weight.rowwise_data.view_as(weights).copy_(weights) + else: + for i in range(num_gemms): + getattr(module, f"weight{i}").copy_(weights[i]) + if use_bias: + if single_grouped_bias: + module.bias.rowwise_data.view_as(biases).copy_(biases) + else: + for i in range(num_gemms): + getattr(module, f"bias{i}").copy_(biases[i]) + + flat_main_grad = None + initial_main_grad = None + weight_params = _grouped_linear_weight_params(module) + if fuse_wgrad_accumulation: + # MCore owns one flat FP32 grad buffer. Discrete parameters receive per-expert + # views, while a single grouped parameter receives one view over the full range. + flat_main_grad = torch.full( + (weights.numel(),), + 0.25, + dtype=torch.float32, + device="cuda", + ) + packed_main_grad = flat_main_grad.view_as(weights) + main_grad_views = ( + [packed_main_grad] + if single_grouped_weight + else [packed_main_grad[i] for i in range(num_gemms)] + ) + for param, main_grad in zip(weight_params, main_grad_views): + param.main_grad = main_grad + param.overwrite_main_grad = False + param.zero_out_wgrad = False + param.grad_added_to_main_grad = False + assert ( + param.main_grad.untyped_storage().data_ptr() + == flat_main_grad.untyped_storage().data_ptr() + ) + initial_main_grad = flat_main_grad.clone() + + x = x_base.detach().clone().requires_grad_(True) + m_splits_arg = ( + torch.tensor(m_splits, dtype=torch.int64, device="cuda") if use_grouped_tensor else m_splits + ) + with autocast(enabled=fp8_recipe is not None, recipe=fp8_recipe): + y = module( + x, + m_splits_arg, + is_first_microbatch=False if fuse_wgrad_accumulation else None, + ) + y.backward(dy) + + if fuse_wgrad_accumulation and delay_wgrad_compute: + torch.testing.assert_close(flat_main_grad, initial_main_grad, rtol=0, atol=0) + + # The grouped-tensor path computes dbias during the main backward even when dW is delayed. + if use_bias and use_grouped_tensor: + assert all(param.grad is not None for param in _grouped_linear_bias_params(module)) + + if delay_wgrad_compute: + module.backward_dw() + + if fuse_wgrad_accumulation: + assert not torch.equal(flat_main_grad, initial_main_grad) + for param in weight_params: + assert param.grad_added_to_main_grad + assert ( + param.main_grad.untyped_storage().data_ptr() + == flat_main_grad.untyped_storage().data_ptr() + ) + packed_wgrad = flat_main_grad.view_as(weights) + elif single_grouped_weight: + packed_wgrad = module.weight.grad.view_as(weights) + else: + packed_wgrad = torch.stack([param.grad for param in weight_params]) + + packed_dbias = None + if use_bias: + bias_params = _grouped_linear_bias_params(module) + if single_grouped_bias: + packed_dbias = bias_params[0].grad.view_as(biases) + else: + packed_dbias = torch.stack([param.grad for param in bias_params]) + + return { + "output": y.detach().clone(), + "dgrad": x.grad.detach().clone(), + "wgrad": packed_wgrad.detach().clone(), + "dbias": None if packed_dbias is None else packed_dbias.detach().clone(), + } + + +_GROUPED_PARAMETER_LAYOUTS = [ + pytest.param(False, False, False, id="no-bias-discrete-weight"), + pytest.param(False, True, False, id="no-bias-single-weight"), + pytest.param(True, False, False, id="bias-discrete-weight-discrete-bias"), + pytest.param(True, True, False, id="bias-single-weight-discrete-bias"), + pytest.param(True, False, True, id="bias-discrete-weight-single-bias"), + pytest.param(True, True, True, id="bias-single-weight-single-bias"), +] + + +@pytest.mark.parametrize( + "use_bias,single_grouped_weight,single_grouped_bias", _GROUPED_PARAMETER_LAYOUTS +) +@pytest.mark.parametrize( + "fp8_recipe", + [ + pytest.param(None, id="bf16"), + pytest.param( + recipe.Float8CurrentScaling(), + marks=pytest.mark.skipif(not _fp8_available, reason=_reason_for_no_fp8), + id="fp8-current-scaling", + ), + pytest.param( + recipe.MXFP8BlockScaling(), + marks=pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8), + id="mxfp8", + ), + pytest.param( + recipe.Float8BlockScaling(), + marks=pytest.mark.skipif( + not _fp8_block_scaling_available, + reason=_reason_for_no_fp8_block_scaling, + ), + id="fp8-block-scaling", + ), + ], +) +@pytest.mark.parametrize("delay_wgrad_compute", _ALL_BOOLEAN) +@pytest.mark.parametrize("fuse_wgrad_accumulation", _ALL_BOOLEAN) +def test_grouped_parameter_layout_matches_cpu_m_splits( + monkeypatch, + use_bias, + single_grouped_weight, + single_grouped_bias, + fp8_recipe, + delay_wgrad_compute, + fuse_wgrad_accumulation, +): + """Match CUDA m_splits and all meaningful parameter layouts against the legacy path.""" + if not is_module_grouped_tensor_path_supported( + fp8_recipe, + torch.bfloat16, + ): + pytest.skip("Recipe is not supported by the module GroupedTensor path on this system.") + FP8GlobalStateManager.reset() + + torch.manual_seed(1234) + # MXFP8 and the grouped FP8 recipes require aligned expert problems. Use the same + # 256-aligned shapes for every recipe so all precision modes exercise one layout. + num_gemms = 2 + in_features = 256 + out_features = 256 + m_splits = [256, 512] + total_tokens = sum(m_splits) + x_base = (0.1 * torch.randn(total_tokens, in_features, device="cuda")).to(torch.bfloat16) + dy = (0.1 * torch.randn(total_tokens, out_features, device="cuda")).to(torch.bfloat16) + weights = (0.1 * torch.randn(num_gemms, out_features, in_features, device="cuda")).to( + torch.bfloat16 + ) + biases = None + if use_bias: + biases = (0.1 * torch.randn(num_gemms, out_features, device="cuda")).to(torch.bfloat16) + + # The CPU m_splits baseline is explicitly the legacy, discrete-parameter contract. + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") + reference = _run_grouped_parameter_layout( + use_grouped_tensor=False, + fp8_recipe=fp8_recipe, + single_grouped_weight=False, + single_grouped_bias=False, + use_bias=use_bias, + delay_wgrad_compute=delay_wgrad_compute, + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + x_base=x_base, + dy=dy, + weights=weights, + biases=biases, + m_splits=m_splits, + ) + + # Enable single parameters only for the CUDA m_splits target. The explicit layout flags + # below still decide whether this particular case uses discrete or grouped parameters. + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + result = _run_grouped_parameter_layout( + use_grouped_tensor=True, + fp8_recipe=fp8_recipe, + single_grouped_weight=single_grouped_weight, + single_grouped_bias=single_grouped_bias, + use_bias=use_bias, + delay_wgrad_compute=delay_wgrad_compute, + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + x_base=x_base, + dy=dy, + weights=weights, + biases=biases, + m_splits=m_splits, + ) + + tolerances = dict(rtol=1e-2, atol=5e-3) + + for name in ("output", "dgrad", "wgrad", "dbias"): + if reference[name] is None: + assert result[name] is None + else: + torch.testing.assert_close( + result[name].float(), + reference[name].float(), + **tolerances, + msg=f"Mismatch for {name}", + ) + + +@pytest.mark.parametrize( + "fp8_recipe", + [ + pytest.param(None, id="bf16"), + pytest.param( + recipe.Float8CurrentScaling(), + marks=pytest.mark.skipif(not _fp8_available, reason=_reason_for_no_fp8), + id="fp8-current-scaling", + ), + pytest.param( + recipe.MXFP8BlockScaling(), + marks=pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8), + id="mxfp8", + ), + pytest.param( + recipe.Float8BlockScaling(), + marks=pytest.mark.skipif( + not _fp8_block_scaling_available, + reason=_reason_for_no_fp8_block_scaling, + ), + id="fp8-block-scaling", + ), + pytest.param( + recipe.NVFP4BlockScaling(disable_stochastic_rounding=True), + marks=pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4), + id="nvfp4", + ), + ], +) +@pytest.mark.parametrize("single_grouped_weight", _ALL_BOOLEAN) +def test_grouped_tensor_save_original_input_matches_saved_grouped_input( + monkeypatch, + fp8_recipe, + single_grouped_weight, +): + """Saving raw input must preserve native grouped forward, dgrad, and wgrad numerics.""" + if not is_module_grouped_tensor_path_supported(fp8_recipe, torch.bfloat16): + pytest.skip("Recipe is not supported by the module GroupedTensor path on this system.") + if single_grouped_weight and fp8_recipe is not None and fp8_recipe.nvfp4(): + pytest.skip( + "NVFP4 grouped GEMM with single_grouped_weight is not supported yet; " + "only discrete weights are supported." + ) + + def reject_split_fallback(*_args, **_kwargs): + pytest.fail("save_original_input unexpectedly selected the split-quantize path") + + monkeypatch.setattr( + "transformer_engine.pytorch.module.grouped_linear._split_quantize", + reject_split_fallback, + ) + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + + torch.manual_seed(1234) + num_gemms = 2 + in_features = 256 + out_features = 256 + m_splits = [256, 512] + total_tokens = sum(m_splits) + x_base = (0.1 * torch.randn(total_tokens, in_features, device="cuda")).to(torch.bfloat16) + dy = (0.1 * torch.randn(total_tokens, out_features, device="cuda")).to(torch.bfloat16) + weights = (0.1 * torch.randn(num_gemms, out_features, in_features, device="cuda")).to( + torch.bfloat16 + ) + + saved_grouped = _run_grouped_parameter_layout( + use_grouped_tensor=True, + fp8_recipe=fp8_recipe, + single_grouped_weight=single_grouped_weight, + single_grouped_bias=False, + use_bias=False, + delay_wgrad_compute=False, + fuse_wgrad_accumulation=False, + x_base=x_base, + dy=dy, + weights=weights, + biases=None, + m_splits=m_splits, + save_original_input=False, + ) + saved_original = _run_grouped_parameter_layout( + use_grouped_tensor=True, + fp8_recipe=fp8_recipe, + single_grouped_weight=single_grouped_weight, + single_grouped_bias=False, + use_bias=False, + delay_wgrad_compute=False, + fuse_wgrad_accumulation=False, + x_base=x_base, + dy=dy, + weights=weights, + biases=None, + m_splits=m_splits, + save_original_input=True, + ) + + for name in ("output", "dgrad", "wgrad"): + torch.testing.assert_close( + saved_original[name].float(), + saved_grouped[name].float(), + rtol=1e-2, + atol=5e-3, + msg=f"Mismatch for {name}", + ) + + def _run_grouped_linear_path( *, enable_grouped_tensor_path: bool, @@ -1839,12 +2650,29 @@ def _run_grouped_linear_path( recipe.MXFP8BlockScaling(), marks=pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8), ), + pytest.param( + recipe.MXFP8BlockScaling(enable_2d_quantization=True), + marks=pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8), + ), pytest.param( recipe.NVFP4BlockScaling(disable_stochastic_rounding=True), marks=pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4), ), + pytest.param( + recipe.Float8BlockScaling(), + marks=pytest.mark.skipif( + not _fp8_block_scaling_available, reason=_reason_for_no_fp8_block_scaling + ), + ), + ], + ids=[ + "bf16", + "fp8_current_scaling", + "mxfp8", + "mxfp8_2d", + "nvfp4", + "fp8_block_scaling", ], - ids=["bf16", "fp8_current_scaling", "mxfp8", "nvfp4"], ) @pytest.mark.parametrize("bias", _ALL_BOOLEAN) @pytest.mark.parametrize("fp8_model_params", _ALL_BOOLEAN) @@ -1853,32 +2681,18 @@ def test_grouped_linear_grouped_tensor_path_matches_legacy( fp8_recipe, bias, fp8_model_params, delay_wgrad_compute, monkeypatch ): use_fp8 = fp8_recipe is not None - device_capability = torch.cuda.get_device_capability() - if not (9, 0) <= device_capability <= (11, 0): - pytest.skip( - "GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell (SM10x and SM110)." - ) if IS_HIP_EXTENSION: pytest.skip("GroupedTensor grouped GEMM needs nvte_grouped_gemm, unsupported on ROCm.") - # MXFP8/NVFP4 grouped quantization kernels require Blackwell, but FP8 per-tensor - # current scaling also runs on the Hopper grouped GEMM path. - is_current_scaling = use_fp8 and fp8_recipe.float8_current_scaling() - if use_fp8 and not is_current_scaling and device_capability < (10, 0): - pytest.skip( - "Quantized GroupedTensor grouped GEMM path (MXFP8/NVFP4) requires Blackwell (SM100+)." - ) - cublaslt_version = tex.get_cublasLt_version() - if device_capability < (10, 0) and cublaslt_version < 130400: - pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.") - if is_current_scaling and device_capability < (10, 0) and cublaslt_version < 130500: - pytest.skip("FP8 per-tensor scaling grouped GEMM on Hopper requires cuBLAS 13.5+.") - if cublaslt_version < 130300: - pytest.skip("Grouped GEMM requires cuBLAS 13.3+.") + dtype = torch.bfloat16 + if not is_module_grouped_tensor_path_supported( + fp8_recipe, + dtype, + ): + pytest.skip("Recipe is not supported by the module GroupedTensor path on this system.") if fp8_model_params and not use_fp8: pytest.skip("fp8_model_params requires FP8") - dtype = torch.bfloat16 num_gemms = 3 in_features = 128 out_features = 128 @@ -1935,8 +2749,11 @@ def test_grouped_linear_grouped_tensor_path_matches_legacy( def test_grouped_linear_grouped_tensor_path_single_grouped_bias_delay_wgrad(monkeypatch): - if torch.cuda.get_device_capability() < (10, 0): - pytest.skip("GroupedTensor grouped GEMM path requires SM100+") + if not is_module_grouped_tensor_path_supported( + None, + torch.bfloat16, + ): + pytest.skip("BF16 GroupedTensor path is unavailable on this system.") monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1") @@ -1965,6 +2782,102 @@ def test_grouped_linear_grouped_tensor_path_single_grouped_bias_delay_wgrad(monk grouped_linear.backward_dw() +def test_grouped_linear_returns_single_grouped_bias_parameter(monkeypatch): + """return_bias preserves the grouped parent and accumulates dbias into it. + + This mirrors how MCore applies a returned MoE bias:: + + x -> GroupedLinear (bias not applied) -> output + + + grouped bias [2, 128] + | + +-> repeat_interleave([256, 256]) + -> per-token bias [512, 128] + | + * routing probabilities + | + +-> loss.backward() -> grouped bias.grad + + Since the loss sums every output feature, each feature of expert ``i`` receives + ``sum(probs_for_expert_i)``. The identity assertion also ensures that TE returns the + registered grouped parent rather than copied or split bias tensors. + """ + if not is_module_grouped_tensor_path_supported( + None, + torch.bfloat16, + ): + pytest.skip("BF16 GroupedTensor path is unavailable on this system.") + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + + dtype = torch.bfloat16 + num_gemms = 2 + in_features = 128 + out_features = 128 + m_splits = torch.tensor([256, 256], dtype=torch.int64, device="cuda") + total_tokens = 512 + grouped_linear = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=True, + return_bias=True, + params_dtype=dtype, + device="cuda", + single_grouped_weight=False, + single_grouped_bias=True, + use_grouped_tensor=True, + ) + + x = torch.randn( + total_tokens, + in_features, + dtype=dtype, + device="cuda", + requires_grad=True, + ) + + probs = torch.cat( + ( + torch.full((256,), 0.25, dtype=dtype, device="cuda"), + torch.full((256,), 0.5, dtype=dtype, device="cuda"), + ) + ) + output, returned_bias = grouped_linear(x, m_splits) + + assert returned_bias is grouped_linear.bias + assert returned_bias.shape == (num_gemms, out_features) + + bias_per_token = torch.repeat_interleave( + returned_bias, + m_splits, + dim=0, + output_size=total_tokens, + ) + biased_output = (output + bias_per_token * probs.reshape(-1, 1)).to(output.dtype) + biased_output.sum().backward() + + # For token t and output feature j, the externally applied bias contributes + # bias[expert(t), j] * probs[t] to the summed loss. Therefore every bias feature for + # expert e receives sum(probs[t]) over that expert's token range. m_splits places the + # first 256 tokens on expert 0 and the remaining 256 tokens on expert 1. + expected_dbias = ( + torch.stack( + (probs[:256].float().sum(), probs[256:].float().sum()), + ) + .unsqueeze(-1) + .expand(num_gemms, out_features) + ) + assert grouped_linear.bias.grad is not None + # Megatron applies the packed bias in BF16. Compare its gradient against an independently + # accumulated FP32 reference with a tolerance appropriate for a 256-element BF16 reduction. + torch.testing.assert_close( + grouped_linear.bias.grad.float(), + expected_dbias, + rtol=5e-2, + atol=5e-3, + ) + + @pytest.mark.parametrize("use_fused_path", [False, True], ids=["legacy", "grouped_tensor"]) @pytest.mark.parametrize("supply", ["out", "dgrad_out", "both"]) def test_grouped_linear_caller_output_buffers(use_fused_path, supply, monkeypatch): @@ -1976,17 +2889,11 @@ def test_grouped_linear_caller_output_buffers(use_fused_path, supply, monkeypatc if use_fused_path: if IS_HIP_EXTENSION: pytest.skip("GroupedTensor grouped GEMM needs nvte_grouped_gemm, unsupported on ROCm.") - device_capability = torch.cuda.get_device_capability() - if not (9, 0) <= device_capability <= (11, 0): - pytest.skip( - "GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell" - " (SM10x and SM110)." - ) - cublaslt_version = tex.get_cublasLt_version() - if device_capability < (10, 0) and cublaslt_version < 130400: - pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.") - if cublaslt_version < 130300: - pytest.skip("Grouped GEMM requires cuBLAS 13.3+.") + if not is_module_grouped_tensor_path_supported( + None, + torch.bfloat16, + ): + pytest.skip("BF16 GroupedTensor path is unavailable on this system.") monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1" if use_fused_path else "0") give_out = supply in ("out", "both") @@ -2064,19 +2971,16 @@ def test_grouped_linear_caller_output_buffers(use_fused_path, supply, monkeypatc @pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4) -def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): - """Non-RHT NVFP4 falls back to the legacy path; check it stays numerically correct. - - Graph-safe grouped quantization currently requires RHT, so requesting NVFP4 with - ``disable_rht=True`` while the fused grouped-tensor path is enabled falls back to the - legacy path internally. We verify the output and gradients against a reference built from - per-GEMM ``te.Linear`` modules that share the same weights and use the same NVFP4 recipe; - the grouped GEMM should match the loop of single GEMMs. +def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(): + """Non-RHT NVFP4 falls back to split-quantize for discrete parameters. + + Graph-safe grouped quantization currently requires RHT. Discrete parameters have a valid + split-quantize fallback, so enabling the grouped-tensor path is a preference rather than a + hard requirement for this parameter layout. """ if torch.cuda.get_device_capability() < (10, 0): pytest.skip("NVFP4 GroupedTensor grouped GEMM path requires SM100+") - monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1") FP8GlobalStateManager.reset() dtype = torch.bfloat16 @@ -2099,7 +3003,6 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): disable_stochastic_rounding=True, ) - # Grouped path: fused path enabled, but non-RHT NVFP4 falls back to legacy internally. grouped_linear = GroupedLinear( num_gemms, in_features, @@ -2107,6 +3010,7 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): bias=False, params_dtype=dtype, device="cuda", + use_grouped_tensor=True, ) with torch.no_grad(): for i in range(num_gemms): @@ -2117,7 +3021,6 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): y = grouped_linear(x, m_splits) y.backward(dy) - # Reference: one te.Linear per GEMM sharing the same weights and NVFP4 recipe. ref_linears = torch.nn.ModuleList( [ Linear(in_features, out_features, bias=False, params_dtype=dtype, device="cuda") @@ -2135,7 +3038,6 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): ) y_ref.backward(dy) - # cuBLAS grouped GEMM should match the loop of single GEMMs bit-for-bit. tols = dict(rtol=0, atol=0) torch.testing.assert_close(y.float(), y_ref.float(), **tols) torch.testing.assert_close(x.grad.float(), x_ref.grad.float(), **tols) @@ -2147,6 +3049,33 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): ) +def test_grouped_linear_delay_wgrad_rejects_implicit_fallback(monkeypatch): + """Delayed wgrad reports when a grouped-tensor request used the legacy path.""" + monkeypatch.setattr( + "transformer_engine.pytorch.module.grouped_linear.is_module_grouped_tensor_path_supported", + lambda *_args, **_kwargs: False, + ) + grouped_linear = GroupedLinear( + 2, + 64, + 64, + bias=True, + params_dtype=torch.bfloat16, + device="cuda", + delay_wgrad_compute=True, + use_grouped_tensor=True, + ) + x = torch.randn(16, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True) + m_splits = torch.tensor([8, 8], dtype=torch.int64, device="cuda") + + grouped_linear(x, m_splits).sum().backward() + with pytest.raises( + RuntimeError, + match="implicit fallback is unsupported with delay_wgrad_compute=True", + ): + grouped_linear.backward_dw() + + @pytest.mark.parametrize( "fp8_recipe", [ @@ -2163,39 +3092,31 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): recipe.NVFP4BlockScaling(disable_stochastic_rounding=True), marks=pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4), ), + pytest.param( + recipe.Float8BlockScaling(), + marks=pytest.mark.skipif( + not _fp8_block_scaling_available, reason=_reason_for_no_fp8_block_scaling + ), + ), ], - ids=["bf16", "fp8_current_scaling", "mxfp8", "nvfp4"], + ids=["bf16", "fp8_current_scaling", "mxfp8", "nvfp4", "fp8_block_scaling"], ) @pytest.mark.parametrize("bias", _ALL_BOOLEAN) def test_grouped_linear_fused_path_cuda_graph_safe(fp8_recipe, bias, monkeypatch): """Fused GroupedTensor GEMM path should be CUDA graph capturable.""" use_fp8 = fp8_recipe is not None - device_capability = torch.cuda.get_device_capability() - if not (9, 0) <= device_capability <= (11, 0): - pytest.skip( - "GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell (SM10x and SM110)." - ) if IS_HIP_EXTENSION: pytest.skip("GroupedTensor grouped GEMM needs nvte_grouped_gemm, unsupported on ROCm.") - # MXFP8/NVFP4 grouped quantization kernels require Blackwell, but FP8 per-tensor - # current scaling also runs on the Hopper grouped GEMM path. - is_current_scaling = use_fp8 and fp8_recipe.float8_current_scaling() - if use_fp8 and not is_current_scaling and device_capability < (10, 0): - pytest.skip( - "Quantized GroupedTensor grouped GEMM path (MXFP8/NVFP4) requires Blackwell (SM100+)." - ) - cublaslt_version = tex.get_cublasLt_version() - if device_capability < (10, 0) and cublaslt_version < 130400: - pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.") - if is_current_scaling and device_capability < (10, 0) and cublaslt_version < 130500: - pytest.skip("FP8 per-tensor scaling grouped GEMM on Hopper requires cuBLAS 13.5+.") - if cublaslt_version < 130300: - pytest.skip("Grouped GEMM requires cuBLAS 13.3+.") + dtype = torch.bfloat16 + if not is_module_grouped_tensor_path_supported( + fp8_recipe, + dtype, + ): + pytest.skip("Recipe is not supported by the module GroupedTensor path on this system.") monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1") FP8GlobalStateManager.reset() - dtype = torch.bfloat16 device = "cuda" num_gemms = 3 in_features = 128 @@ -2296,6 +3217,24 @@ def _train_step(x, dy, out_buf, *, use_graphed): torch.testing.assert_close(graph_grad.float(), param.grad.float(), **tols) +@pytest.mark.skipif(not _fp8_block_scaling_available, reason=_reason_for_no_fp8_block_scaling) +@pytest.mark.skipif( + not (10, 0) <= torch.cuda.get_device_capability() <= (11, 0), + reason="Error path only triggers on Blackwell (SM100/SM110).", +) +def test_grouped_linear_fused_path_fp8_block_scaling_blackwell_error(monkeypatch): + """FP8BS + fused env var on Blackwell must raise, not silently fall back.""" + monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1") + FP8GlobalStateManager.reset() + dtype = torch.bfloat16 + grouped_linear = GroupedLinear(2, 128, 128, bias=False, params_dtype=dtype, device="cuda") + x = torch.randn(256, 128, device="cuda", dtype=dtype, requires_grad=True) + m_splits = torch.tensor([128, 128], dtype=torch.int64, device="cuda") + with pytest.raises(RuntimeError, match="Hopper-only"): + with autocast(enabled=True, recipe=recipe.Float8BlockScaling()): + grouped_linear(x, m_splits) + + @pytest.mark.parametrize("swizzle_type", ["mxfp8_rowwise", "mxfp8_columnwise", "nvfp4"]) def test_swizzle_scales_and_pack_ptrs_for_discrete_weights( swizzle_type: str, diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 76550db5f8..c6d73cba63 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -15,6 +15,8 @@ import torch import transformer_engine.pytorch as te +from transformer_engine.pytorch.constants import TE_DType +import transformer_engine.pytorch.ops.fused.grouped_mlp as grouped_mlp_module from transformer_engine.pytorch.ops.fused.grouped_mlp import ( _cudnn_frontend_supports_grouped_gemm_srelu, _cudnn_frontend_version_supported, @@ -22,6 +24,7 @@ from transformer_engine.pytorch.ops.basic.grouped_linear import ( OUTPUT_BUFFER_KEY, GRAD_INPUT_BUFFER_KEY, + is_op_fuser_grouped_tensor_path_supported, ) from transformer_engine.pytorch import ( QuantizedTensor, @@ -229,9 +232,73 @@ def make_reference_and_test_tensors( return ref, test +class _InjectGrad(torch.autograd.Function): + """Replace the gradient flowing into ``x`` with ``grad``. + + Mirrors how an FP8 token dispatch delivers a pre-quantized ``GroupedTensor`` + grad output: the downstream op simply returns one as its grad input. + """ + + @staticmethod + def forward(ctx, x, grad): # pylint: disable=arguments-differ + ctx.injected_grad = grad + return x + + @staticmethod + def backward(ctx, grad_output): # pylint: disable=arguments-differ + return ctx.injected_grad, None + + class TestGroupedLinearOp: """Tests for advanced features with grouped linear basic op""" + def test_meta_single_grouped_weight_with_delayed_wgrad(self, monkeypatch) -> None: + """A deferred op shell must not access its grouped parent before it is attached.""" + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + + op = te.ops.GroupedLinear( + 2, + 16, + 16, + device="meta", + dtype=torch.bfloat16, + single_grouped_weight=True, + delay_wgrad_compute=True, + ) + + assert op._parameters.get("weight") is None + assert op.weight0.device.type == "meta" + + # Mirror an external caller attaching the grouped parent after constructing + # the parameterless shell, then verify delayed-wgrad metadata is applied. + grouped_weight = torch.nn.Parameter(torch.empty(2, 16, 16, device="meta")) + assert not hasattr(grouped_weight, "skip_backward_post_hook") + op.register_parameter("weight", grouped_weight) + for group_idx in range(op.num_groups): + op.register_parameter(f"weight{group_idx}", None) + + assert grouped_weight.skip_backward_post_hook + + def test_single_grouped_bias_uses_registered_packed_storage(self, monkeypatch) -> None: + """The grouped bias compute view must alias the registered trainable parent.""" + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + op = te.ops.GroupedLinear( + 2, + 128, + 128, + bias=True, + device="cuda", + dtype=torch.bfloat16, + single_grouped_bias=True, + ) + + bias_packed = op._get_packed_bias_tensor(torch.bfloat16) + + assert bias_packed.shape == (op.num_groups, op.out_features) + assert bias_packed.untyped_storage().data_ptr() == op.bias.rowwise_data.data_ptr() + assert op.bias.requires_grad + assert dict(op.named_parameters())["bias"] is op.bias + @pytest.mark.parametrize("bias", (False, True)) @pytest.mark.parametrize("dtype", _dtypes) @pytest.mark.parametrize("quantization", _quantization_list) @@ -292,6 +359,16 @@ def test_grouped_linear( if single_grouped_bias and not bias: pytest.skip("single_grouped_bias requires bias=True") + recipe = make_recipe(quantization) + compute_recipe = recipe if quantized_compute else None + if ( + single_grouped_weight or single_grouped_bias + ) and not is_op_fuser_grouped_tensor_path_supported( + compute_recipe, + dtype, + ): + # Single grouped parameters intentionally have no split-quantize fallback. + pytest.skip("Single grouped parameters require the native grouped-tensor path") if single_grouped_weight and quantized_weight and quantization in ("fp8_delayed_scaling"): pytest.skip( "single_grouped_weight does not support FP8 delayed scaling " @@ -347,7 +424,6 @@ def test_grouped_linear( y_ref.backward(dy_ref) # Construct fusible operation - recipe = make_recipe(quantization) with te.quantized_model_init(enabled=quantized_weight, recipe=recipe): op = te.ops.GroupedLinear( group_size, @@ -436,6 +512,189 @@ def test_grouped_linear( else: assert b_test.grad is None + @staticmethod + def _make_rowwise_mxfp8_wire_input( + x_hp: torch.Tensor, + group_size: int, + split_sizes: torch.Tensor, + ) -> "GroupedTensor": + """Rowwise-only MXFP8 GroupedTensor with compact scales (FP8 dispatch wire format).""" + wire_quantizer = MXFP8Quantizer( + fp8_dtype=TE_DType[torch.float8_e4m3fn], rowwise=True, columnwise=False + ) + x_wire = tex.group_quantize( + x_hp, wire_quantizer, group_size, split_sizes.to(dtype=torch.int64) + ) + # Sanity: the wire tensor is rowwise-only with unswizzled scales. + assert x_wire.columnwise_data is None + assert not x_wire._with_gemm_swizzled_scales + return x_wire + + @pytest.mark.parametrize("weight_requires_grad", (False, True)) + def test_grouped_linear_prequantized_mxfp8_input( + self, + *, + group_size: int = 4, + weight_shape: tuple[int, int] = (256, 256), + split_alignment: int = 128, + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + weight_requires_grad: bool, + ) -> None: + """Rowwise-only MXFP8 GroupedTensor input (FP8 token dispatch wire format). + + The input arrives already rowwise-quantized with compact scales. The op + must feed the rowwise data to the forward GEMM as-is and manufacture the + columnwise copy for the wgrad GEMM. The reference run consumes the + *dequantized* wire tensor (the only data a layer can see after FP8 + dispatch) on the normal quantize-from-BF16 path; because MXFP8 requant + is idempotent along the rowwise axis and both paths derive the + columnwise copy from the same dequantized data, the two runs must match + bit-for-bit. + """ + if not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + maybe_skip_quantization("mxfp8", dims=weight_shape, device=device, dtype=dtype) + + # Split sizes (including an empty group) + split_sizes = [split_alignment * i for i in range(group_size)] + random.shuffle(split_sizes) + split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device) + + out_features, in_features = weight_shape + total_tokens = int(split_sizes.sum().item()) + in_shape = (total_tokens, in_features) + + # Wire-format input and its exact dequantization (the reference input). + x_hp = torch.rand(in_shape, dtype=dtype, device=device) - 0.5 + x_wire = self._make_rowwise_mxfp8_wire_input(x_hp, group_size, split_sizes) + x_ref = tex.group_dequantize(x_wire, TE_DType[dtype]).rowwise_data.view(in_shape) + + dy = torch.rand((total_tokens, out_features), dtype=dtype, device=device) - 0.5 + recipe = make_recipe("mxfp8") + op = te.ops.GroupedLinear( + group_size, in_features, out_features, bias=False, device=device, dtype=dtype + ) + with torch.no_grad(): + for param in op.parameters(): + param.requires_grad_(requires_grad=weight_requires_grad) + + def _run(x): + with te.autocast(enabled=True, recipe=recipe): + y = op(x, split_sizes) + wgrads = [] + if weight_requires_grad: + y.backward(dy) + for group_idx in range(group_size): + weight = getattr(op, f"weight{group_idx}") + wgrads.append(weight.grad.detach().clone()) + weight.grad = None + return y.detach(), wgrads + + y_ref, wgrads_ref = _run(x_ref) + y_test, wgrads_test = _run(x_wire) + + # Bit-exact match expected (identical quantized inputs and kernels). + torch.testing.assert_close(y_test, y_ref, rtol=0, atol=0) + for wgrad_test, wgrad_ref in zip(wgrads_test, wgrads_ref): + torch.testing.assert_close(wgrad_test, wgrad_ref, rtol=0, atol=0) + + @pytest.mark.parametrize("bias_mode", ("none", "plain", "scaled")) + @pytest.mark.parametrize("weight_requires_grad", (False, True)) + def test_grouped_linear_prequantized_mxfp8_grad( + self, + *, + group_size: int = 4, + weight_shape: tuple[int, int] = (256, 256), + split_alignment: int = 128, + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + bias_mode: str, + weight_requires_grad: bool, + ) -> None: + """Rowwise-only MXFP8 GroupedTensor grad output (FP8 token dispatch, backward). + + Mirrors ``test_grouped_linear_prequantized_mxfp8_input`` on the backward + side: the rowwise data feeds the dgrad GEMM and the columnwise copy is + manufactured for wgrad. Covers all three bias gradient sources: none, + ``plain`` (fused into the columnwise stage of the quantize kernel, or + reduced from the dequantized grad when frozen weights leave no columnwise + stage), and ``scaled`` (``scale_bias``, whose dbias/dscales need the + dequantized grad because they depend on the routing probabilities). + + With frozen weights TE also requires frozen biases, so bias gradients are + not observable there; ``dscales`` still is, since the probabilities are an + input rather than a parameter. + """ + if not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + maybe_skip_quantization("mxfp8", dims=weight_shape, device=device, dtype=dtype) + + has_bias = bias_mode != "none" + use_scale_bias = bias_mode == "scaled" + + # Split sizes (including an empty group) + split_sizes = [split_alignment * i for i in range(group_size)] + random.shuffle(split_sizes) + split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device) + + out_features, in_features = weight_shape + total_tokens = int(split_sizes.sum().item()) + + x = torch.rand((total_tokens, in_features), dtype=dtype, device=device) - 0.5 + dy_hp = torch.rand((total_tokens, out_features), dtype=dtype, device=device) - 0.5 + probs = torch.rand((total_tokens,), dtype=dtype, device=device) + + # Wire-format grad output and its exact dequantization (the reference grad). + dy_wire = self._make_rowwise_mxfp8_wire_input(dy_hp, group_size, split_sizes) + dy_ref = tex.group_dequantize(dy_wire, TE_DType[dtype]).rowwise_data.view( + total_tokens, out_features + ) + + recipe = make_recipe("mxfp8") + op = te.ops.GroupedLinear( + group_size, + in_features, + out_features, + bias=has_bias, + device=device, + dtype=dtype, + scale_bias=use_scale_bias, + ) + if not weight_requires_grad: + # TE requires bias.requires_grad to match weight.requires_grad. + for param in op.parameters(): + param.requires_grad_(False) + + def _run(grad): + x_in = x.detach().clone().requires_grad_() + probs_in = probs.detach().clone().requires_grad_(use_scale_bias) + extra_inputs = (split_sizes, probs_in) if use_scale_bias else (split_sizes,) + with te.autocast(enabled=True, recipe=recipe): + y = op(x_in, *extra_inputs) + # Deliver ``grad`` as the op's grad output, as FP8 dispatch would. + _InjectGrad.apply(y, grad).backward(torch.ones_like(y)) + grads = [("dx", x_in.grad)] + if use_scale_bias: + grads.append(("dprobs", probs_in.grad)) + if weight_requires_grad: + for group_idx in range(group_size): + weight = getattr(op, f"weight{group_idx}") + grads.append((f"w{group_idx}", weight.grad.detach().clone())) + weight.grad = None + if has_bias: + bias_param = getattr(op, f"bias{group_idx}") + grads.append((f"b{group_idx}", bias_param.grad.detach().clone())) + bias_param.grad = None + return grads + + grads_ref = _run(dy_ref) + grads_test = _run(dy_wire) + + # Bit-exact match expected (identical quantized grads and kernels). + for (_, grad_test), (_, grad_ref) in zip(grads_test, grads_ref): + torch.testing.assert_close(grad_test, grad_ref, rtol=0, atol=0) + @pytest.mark.parametrize("dtype", (torch.bfloat16, torch.float16)) @pytest.mark.parametrize( "quantization", @@ -478,22 +737,6 @@ def test_grouped_linear_cuda_graph_safe( "single_grouped_weight/single_grouped_bias requires" " NVTE_GROUPED_LINEAR_SINGLE_PARAM=1" ) - device_capability = torch.cuda.get_device_capability() - if device_capability < (9, 0): - pytest.skip( - "Grouped GEMM CUDA-graph-safe path requires Hopper (SM90) or Blackwell (SM100+)" - ) - # BF16/FP16 and FP8 per-tensor current scaling run on the Hopper grouped GEMM path, - # but MXFP8/NVFP4 grouped quantization kernels require Blackwell (SM100+). - requires_blackwell = quantization is not None and quantization != "fp8_current_scaling" - if requires_blackwell and device_capability < (10, 0): - pytest.skip("MXFP8/NVFP4 grouped GEMM CUDA-graph-safe path requires SM100+ (Blackwell)") - # Grouped GEMM on Hopper requires cuBLAS 13.4+; Blackwell requires cuBLAS 13.3+. - cublaslt_version = tex.get_cublasLt_version() - if device_capability < (10, 0) and cublaslt_version < 130400: - pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.") - if cublaslt_version < 130300: - pytest.skip("Grouped GEMM requires cuBLAS 13.3+.") if quantization is None and quantized_weight: pytest.skip("quantized_weight requires a quantization recipe") if ( @@ -511,6 +754,13 @@ def test_grouped_linear_cuda_graph_safe( "only discrete weights (single_grouped_weight=False) are supported." ) + recipe = make_recipe(quantization) + if not is_op_fuser_grouped_tensor_path_supported( + recipe, + dtype, + ): + pytest.skip("Configuration falls back to the non-CUDA-graph-safe path") + single_grouped_bias = bias and single_grouped_weight # Split sizes (statically pinned for graph capture) @@ -523,7 +773,6 @@ def test_grouped_linear_cuda_graph_safe( in_shape = (num_active_tokens + token_padding, in_features) out_shape = (in_shape[0], out_features) - recipe = make_recipe(quantization) with te.quantized_model_init(enabled=quantized_weight, recipe=recipe): op = te.ops.GroupedLinear( group_size, @@ -701,7 +950,10 @@ def test_grouped_mlp( """GroupedLinear + scaled activation + GroupedLinear""" # Split sizes - split_sizes = [split_alignment * (i) for i in range(group_size)] + if group_size == 1: + split_sizes = [split_alignment] + else: + split_sizes = [split_alignment * i for i in range(group_size)] random.shuffle(split_sizes) split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device) @@ -724,6 +976,13 @@ def test_grouped_mlp( maybe_skip_quantization(quantization, dims=in_shape, device=device, dtype=dtype) if dtype == torch.bfloat16 and not is_bf16_available(): pytest.skip("BF16 requires SM 8.0+") + if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and ( + single_grouped_weight or single_grouped_bias + ): + pytest.skip( + "single_grouped_weight/single_grouped_bias requires" + " NVTE_GROUPED_LINEAR_SINGLE_PARAM=1" + ) if single_grouped_weight and quantization != "mxfp8": pytest.skip("single_grouped_weight is only supported for MXFP8 quantization") if single_grouped_bias and not bias: @@ -995,8 +1254,10 @@ def _make_module(): or ( quantization == "nvfp4_rht" and dtype == torch.bfloat16 - and activation == "scaled_srelu" - and glu_interleave_size is None + and ( + (not activation_is_glu and glu_interleave_size is None) + or (activation_is_glu and glu_interleave_size == 32) + ) ) ) if expected_grouped_mlp_fusion: @@ -1076,6 +1337,208 @@ def _make_module(): assert_close(fc1.weight.grad, fc1_w_ref_grad, **tols) assert_close(fc2.weight.grad, fc2_w_ref_grad, **tols) + @pytest.mark.parametrize("weight_requires_grad", (False, True)) + def test_grouped_mlp_prequantized_mxfp8_input( + self, + *, + group_size: int = 4, + hidden_size: int = 256, + split_alignment: int = 256, + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + weight_requires_grad: bool, + ) -> None: + """Fused grouped MLP with a rowwise-only MXFP8 GroupedTensor input. + + Production (FP8 token dispatch) path: FC1 receives an already + rowwise-quantized input with compact scales. The fused op must feed the + rowwise data to the forward GEMM and manufacture FC1's columnwise copy + for wgrad. Compared bit-for-bit against a run on the dequantized wire + input (see ``test_grouped_linear_prequantized_mxfp8_input``). + """ + if not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): + pytest.skip("Fused grouped MLP (CuTeDSL) is not supported on this system") + maybe_skip_quantization( + "mxfp8", dims=(hidden_size, hidden_size), device=device, dtype=dtype + ) + + # Split sizes (including an empty group); sum is a multiple of 128. + split_sizes = [split_alignment * i for i in range(group_size)] + random.shuffle(split_sizes) + split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device) + total_tokens = int(split_sizes.sum().item()) + glu_interleave_size = 32 + + # Wire-format FC1 input and its exact dequantization (reference input). + x_hp = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + x_wire = TestGroupedLinearOp._make_rowwise_mxfp8_wire_input(x_hp, group_size, split_sizes) + x_ref = tex.group_dequantize(x_wire, TE_DType[dtype]).rowwise_data.view( + total_tokens, hidden_size + ) + + probs = torch.rand((total_tokens,), dtype=dtype, device=device) + dy = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + + recipe = make_recipe("mxfp8") + with te.quantized_model_init(enabled=True, recipe=recipe): + fc1 = te.ops.GroupedLinear( + group_size, hidden_size, 2 * hidden_size, bias=False, device=device, dtype=dtype + ) + fc2 = te.ops.GroupedLinear( + group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype + ) + module = te.ops.Sequential( + fc1, te.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size), fc2 + ) + with torch.no_grad(): + for param in module.parameters(): + param.requires_grad_(requires_grad=weight_requires_grad) + + def _run(x): + with te.autocast(enabled=True, recipe=recipe): + y = module(x, split_sizes, probs, split_sizes) + fc1_wgrads, fc2_wgrads = [], [] + if weight_requires_grad: + y.backward(dy) + for group_idx in range(group_size): + fc1_w = getattr(fc1, f"weight{group_idx}") + fc2_w = getattr(fc2, f"weight{group_idx}") + fc1_wgrads.append(fc1_w.grad.detach().clone()) + fc2_wgrads.append(fc2_w.grad.detach().clone()) + fc1_w.grad = None + fc2_w.grad = None + return y.detach(), fc1_wgrads, fc2_wgrads + + y_ref, fc1_wgrads_ref, fc2_wgrads_ref = _run(x_ref) + y_test, fc1_wgrads_test, fc2_wgrads_test = _run(x_wire) + + # Confirm the CuTeDSL fused op was actually formed (not the fallback). + forward_ops = module._module_groups[0]._forward_ops + assert len(forward_ops) == 1 + assert isinstance(forward_ops[0][0], te.ops.fused.GroupedMLP_CuTeGEMMGLU) + + # Bit-exact match expected (identical quantized inputs and kernels). + torch.testing.assert_close(y_test, y_ref, rtol=0, atol=0) + for wgrad_test, wgrad_ref in zip(fc1_wgrads_test, fc1_wgrads_ref): + torch.testing.assert_close(wgrad_test, wgrad_ref, rtol=0, atol=0) + for wgrad_test, wgrad_ref in zip(fc2_wgrads_test, fc2_wgrads_ref): + torch.testing.assert_close(wgrad_test, wgrad_ref, rtol=0, atol=0) + + @pytest.mark.parametrize("bias", (False, True)) + @pytest.mark.parametrize("weight_requires_grad", (False, True)) + def test_grouped_mlp_prequantized_mxfp8_grad( + self, + *, + group_size: int = 4, + hidden_size: int = 256, + split_alignment: int = 256, + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + bias: bool, + weight_requires_grad: bool, + ) -> None: + """Fused grouped MLP with a rowwise-only MXFP8 GroupedTensor grad output. + + FC2 receives the pre-quantized grad, as an FP8 token dispatch delivers it + on the backward pass. With ``bias`` FC2 uses ``scale_bias``, whose + dbias/dscales need the dequantized grad rather than the fused dbias. + """ + if not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): + pytest.skip("Fused grouped MLP (CuTeDSL) is not supported on this system") + if not weight_requires_grad: + # Independent of pre-quantization: the fused forward saves the input + # activations whenever anything requires grad, then asserts they carry + # columnwise data -- which is only built when the weights need grads. + pytest.skip("Fused grouped MLP does not support frozen weights") + maybe_skip_quantization( + "mxfp8", dims=(hidden_size, hidden_size), device=device, dtype=dtype + ) + + # Split sizes (including an empty group); sum is a multiple of 128. + split_sizes = [split_alignment * i for i in range(group_size)] + random.shuffle(split_sizes) + split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device) + total_tokens = int(split_sizes.sum().item()) + glu_interleave_size = 32 + + x = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + probs = torch.rand((total_tokens,), dtype=dtype, device=device) + dy_hp = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + + # Wire-format grad output and its exact dequantization (the reference grad). + dy_wire = TestGroupedLinearOp._make_rowwise_mxfp8_wire_input(dy_hp, group_size, split_sizes) + dy_ref = tex.group_dequantize(dy_wire, TE_DType[dtype]).rowwise_data.view( + total_tokens, hidden_size + ) + + recipe = make_recipe("mxfp8") + with te.quantized_model_init(enabled=True, recipe=recipe): + fc1 = te.ops.GroupedLinear( + group_size, hidden_size, 2 * hidden_size, bias=bias, device=device, dtype=dtype + ) + fc2 = te.ops.GroupedLinear( + group_size, + hidden_size, + hidden_size, + bias=bias, + device=device, + dtype=dtype, + scale_bias=bias, + ) + module = te.ops.Sequential( + fc1, te.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size), fc2 + ) + + # Frozen experts (weights) with a still-training bias/router is the case + # where the dequantized grad is needed but is not a byproduct of wgrad. + if not weight_requires_grad: + for fc in (fc1, fc2): + for group_idx in range(group_size): + getattr(fc, f"weight{group_idx}").requires_grad_(False) + + def _run(grad): + x_in = x.detach().clone().requires_grad_() + fc2_extra = (split_sizes, probs) if bias else (split_sizes,) + with te.autocast(enabled=True, recipe=recipe): + y = module(x_in, split_sizes, probs, *fc2_extra) + _InjectGrad.apply(y, grad).backward(torch.ones_like(y)) + grads = [("dx", x_in.grad)] + for name, fc in (("fc1", fc1), ("fc2", fc2)): + for group_idx in range(group_size): + if weight_requires_grad: + weight = getattr(fc, f"weight{group_idx}") + grads.append((f"{name}_w{group_idx}", weight.grad.detach().clone())) + weight.grad = None + if bias: + bias_param = getattr(fc, f"bias{group_idx}") + grads.append((f"{name}_b{group_idx}", bias_param.grad.detach().clone())) + bias_param.grad = None + return grads + + grads_ref = _run(dy_ref) + grads_test = _run(dy_wire) + + # Confirm the CuTeDSL fused op was actually formed (not the fallback). + forward_ops = module._module_groups[0]._forward_ops + assert len(forward_ops) == 1 + assert isinstance(forward_ops[0][0], te.ops.fused.GroupedMLP_CuTeGEMMGLU) + + # Bit-exact match expected (identical quantized grads and kernels), except + # bias gradients: the fused kernels generate them with an accumulation + # that is not reproducible run to run (two runs on identical inputs differ + # by one BF16 ulp). Same tolerances as + # ``test_grouped_mlp_single_weight_numerics``. + bias_tols = {"rtol": 0.05, "atol": 0.015625} + for (name, grad_test), (_, grad_ref) in zip(grads_test, grads_ref): + if "_b" in name: + torch.testing.assert_close(grad_test, grad_ref, **bias_tols) + else: + torch.testing.assert_close(grad_test, grad_ref, rtol=0, atol=0) + @pytest.mark.parametrize("bias", (False, True)) @pytest.mark.parametrize("quantization", _grouped_mlp_quantization_list) @pytest.mark.parametrize( @@ -1104,6 +1567,95 @@ def test_grouped_mlp_fp16( activation=activation, ) + @pytest.mark.parametrize("bias", (False, True)) + @pytest.mark.parametrize("runtime_offsets_supported", (False, True)) + def test_grouped_mlp_single_group_mxfp8( + self, + monkeypatch, + *, + bias: bool, + runtime_offsets_supported: bool, + ) -> None: + """Single-group GroupedLinear + ScaledSwiGLU + GroupedLinear with MXFP8.""" + if ( + runtime_offsets_supported + and not grouped_mlp_module._cudnn_frontend_supports_single_group_runtime_offsets( + te.ops.ScaledSwiGLU + ) + ): + pytest.skip("Requires cuDNN frontend >= 1.27.0") + monkeypatch.setattr( + grouped_mlp_module, + "_cudnn_frontend_supports_single_group_runtime_offsets", + lambda _activation_type: runtime_offsets_supported, + ) + self.test_grouped_mlp( + group_size=1, + bias=bias, + hidden_size=128, + quantization="mxfp8", + single_grouped_weight=False, + activation="scaled_swiglu", + ) + + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_single_grouped_weight_eval_preserves_columnwise_usage( + self, + *, + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + group_size: int = 2, + hidden_size: int = 128, + ) -> None: + """Eager eval must not drop storage still used by captured training dgrad.""" + + if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): + pytest.skip("MXFP8 fused grouped MLP is not supported on this system") + + split_sizes = torch.full( + (group_size,), + 128, + dtype=torch.int64, + device=device, + ) + num_tokens = int(split_sizes.sum()) + recipe = make_recipe("mxfp8") + + with te.quantized_model_init(enabled=True, recipe=recipe): + fc1 = te.ops.GroupedLinear( + group_size, + hidden_size, + 2 * hidden_size, + device=device, + dtype=dtype, + single_grouped_weight=True, + ) + fc2 = te.ops.GroupedLinear( + group_size, + hidden_size, + hidden_size, + device=device, + dtype=dtype, + single_grouped_weight=True, + ) + module = te.ops.Sequential( + fc1, + te.ops.ScaledSwiGLU(glu_interleave_size=32), + fc2, + ) + + x = torch.randn(num_tokens, hidden_size, device=device, dtype=dtype, requires_grad=True) + probs = torch.ones(num_tokens, device=device, dtype=dtype) + with te.autocast(enabled=True, recipe=recipe): + module(x, split_sizes, probs, split_sizes) + assert fc1.weight.quantizer.columnwise_usage + assert fc2.weight.quantizer.columnwise_usage + + with torch.no_grad(), te.autocast(enabled=True, recipe=recipe): + module(x.detach(), split_sizes, probs, split_sizes) + assert fc1.weight.quantizer.columnwise_usage + assert fc2.weight.quantizer.columnwise_usage + @pytest.mark.parametrize("quantization", _grouped_mlp_quantization_list) @pytest.mark.parametrize("single_grouped_weight", (False, True)) @pytest.mark.parametrize("accumulate_into_main_grad", (False, True)) @@ -1148,6 +1700,8 @@ def test_grouped_mlp_single_weight_numerics( ) -> None: """single_grouped_weight=True/False should match exactly for fused MXFP8 grouped MLP.""" + if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0": + pytest.skip("single_grouped_weight requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1") if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): pytest.skip("MXFP8 fused grouped MLP is not supported on this system") @@ -1466,6 +2020,8 @@ def test_grouped_mlp_overwrite_main_grad( that read ``.grad`` don't see stale bytes from the cached dummy). """ + if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and single_grouped_weight: + pytest.skip("single_grouped_weight requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1") if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): pytest.skip("MXFP8 fused grouped MLP is not supported on this system") @@ -1597,6 +2153,8 @@ def test_grouped_mlp_cuda_graph_safe_mxfp8( ) -> None: """Grouped MLP forward+backward should be CUDA graph capturable (MXFP8).""" + if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and single_grouped_weight: + pytest.skip("single_grouped_weight requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1") if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): pytest.skip("MXFP8 fused grouped MLP is not supported on this system") if dtype not in (torch.bfloat16, torch.float16): diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index 870461e717..146eb60a8a 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -590,7 +590,7 @@ def test_quantize_grouped_mxfp8(self, shape_case: str, output_dbias: bool) -> No if output_dbias: expected_dbias = torch.stack([t.sum(dim=0) for t in input_tensors]) - assert torch.allclose(dbias, expected_dbias) + torch.testing.assert_close(dbias, expected_dbias, rtol=1e-5, atol=4e-3) @pytest.mark.parametrize("output_dbias", [False, True]) @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) @@ -821,6 +821,152 @@ def _run(inp): assert torch.equal(static_output.rowwise_data, expected.rowwise_data) assert torch.equal(static_output.scale_inv, expected.scale_inv) + @pytest.mark.parametrize( + "quantization", + [ + pytest.param( + "fp8_current_scaling", + marks=pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8), + ), + pytest.param( + "fp8_blockwise", + marks=pytest.mark.skipif( + not fp8_block_scaling_grouped_available, + reason=reason_for_no_fp8_block_scaling_grouped, + ), + ), + ], + ) + def test_group_quantize_reuses_destination_and_honors_noop_fp8(self, quantization: str) -> None: + """Check pointer-stable weight-cache updates and no-op CUDA graph replays.""" + num_tensors = 2 + shape = (num_tensors * 256, 256) + if quantization == "fp8_current_scaling": + quantizer = Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + force_pow_2_scales=False, + amax_epsilon=0.0, + ) + quantizer.set_usage(rowwise=True, columnwise=True) + else: + quantizer = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + force_pow_2_scales=False, + amax_epsilon=0.0, + block_scaling_dim=1, + ) + + source = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + output = tex.group_quantize(source, quantizer, num_tensors, None) + output_buffers = tuple( + buffer + for buffer in ( + output.rowwise_data, + output.columnwise_data, + output.scale_inv, + output.columnwise_scale_inv, + output.amax, + output.columnwise_amax, + output.scale, + ) + if buffer is not None + ) + assert output_buffers + pointers = tuple(buffer.data_ptr() for buffer in output_buffers) + + # Re-quantization must update the existing allocations because a captured GEMM keeps + # their raw addresses rather than looking them up again through the Python object. + updated_source = source + 1 + updated = tex.group_quantize( + updated_source, + quantizer, + num_tensors, + None, + output=output, + ) + reference = tex.group_quantize(updated_source, quantizer, num_tensors, None) + assert updated is output + assert pointers == tuple(buffer.data_ptr() for buffer in output_buffers) + reference_buffers = tuple( + buffer + for buffer in ( + reference.rowwise_data, + reference.columnwise_data, + reference.scale_inv, + reference.columnwise_scale_inv, + reference.amax, + reference.columnwise_amax, + reference.scale, + ) + if buffer is not None + ) + assert len(output_buffers) == len(reference_buffers) + for output_buffer, reference_buffer in zip(output_buffers, reference_buffers, strict=True): + assert torch.equal(output_buffer, reference_buffer) + + # A nonzero device flag must preserve every cached payload and metadata buffer. + old_buffers = tuple(buffer.clone() for buffer in output_buffers) + noop = torch.ones(1, dtype=torch.float32, device="cuda") + tex.group_quantize( + updated_source * 2, + quantizer, + num_tensors, + None, + noop_flag=noop, + output=output, + ) + for output_buffer, old_buffer in zip(output_buffers, old_buffers, strict=True): + assert torch.equal(output_buffer, old_buffer) + + # Capture one pointer-stable update. Replays change only source contents and the device + # no-op value, matching TE's CUDA graph microbatch weight-cache protocol. + graph_source = updated_source.clone() + noop.zero_() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + tex.group_quantize( + graph_source, + quantizer, + num_tensors, + None, + noop_flag=noop, + output=output, + ) + + graph_source.copy_(source - 1) + noop.zero_() + graph.replay() + torch.cuda.synchronize() + reference = tex.group_quantize(graph_source, quantizer, num_tensors, None) + reference_buffers = tuple( + buffer + for buffer in ( + reference.rowwise_data, + reference.columnwise_data, + reference.scale_inv, + reference.columnwise_scale_inv, + reference.amax, + reference.columnwise_amax, + reference.scale, + ) + if buffer is not None + ) + assert len(output_buffers) == len(reference_buffers) + for output_buffer, reference_buffer in zip(output_buffers, reference_buffers, strict=True): + assert torch.equal(output_buffer, reference_buffer) + + old_buffers = tuple(buffer.clone() for buffer in output_buffers) + graph_source.copy_(source + 2) + noop.fill_(1) + graph.replay() + torch.cuda.synchronize() + for output_buffer, old_buffer in zip(output_buffers, old_buffers, strict=True): + assert torch.equal(output_buffer, old_buffer) + @pytest.mark.parametrize("mode", ["rowwise", "columnwise", "both"]) @pytest.mark.parametrize( "shape_case", @@ -1027,11 +1173,19 @@ def _assert_fp8_cs_group_quantize_matches_reference( @pytest.mark.parametrize("shape_case", ["uniform", "varying_first"]) @pytest.mark.parametrize("direction", ["rowwise", "columnwise", "both"]) @pytest.mark.parametrize("output_dbias", [False, True]) + @pytest.mark.parametrize( + "force_pow_2_scales", [False, True], ids=["fp32_scales", "pow2_scales"] + ) @pytest.mark.skipif( not fp8_block_scaling_grouped_available, reason=reason_for_no_fp8_block_scaling_grouped ) def test_quantize_grouped_fp8_blockwise( - self, block_scaling_dim: int, shape_case: str, direction: str, output_dbias: bool + self, + block_scaling_dim: int, + shape_case: str, + direction: str, + output_dbias: bool, + force_pow_2_scales: bool, ) -> None: """Test grouped FP8 block-scaling quantization against per-tensor quantization. @@ -1080,7 +1234,7 @@ def test_quantize_grouped_fp8_blockwise( fp8_dtype=tex.DType.kFloat8E4M3, rowwise=rowwise, columnwise=columnwise, - force_pow_2_scales=False, + force_pow_2_scales=force_pow_2_scales, amax_epsilon=0.0, block_scaling_dim=block_scaling_dim, ) @@ -1100,7 +1254,7 @@ def test_quantize_grouped_fp8_blockwise( fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True, - force_pow_2_scales=False, + force_pow_2_scales=force_pow_2_scales, amax_epsilon=0.0, block_scaling_dim=block_scaling_dim, ) @@ -1132,7 +1286,7 @@ def test_quantize_grouped_fp8_blockwise( if output_dbias: expected_dbias = torch.stack([t.sum(dim=0) for t in input_tensors]) - assert torch.allclose(dbias, expected_dbias) + torch.testing.assert_close(dbias, expected_dbias, rtol=1e-5, atol=4e-3) @pytest.mark.parametrize( "shape", @@ -1215,6 +1369,147 @@ def test_group_dequantize_cudagraph_capturable(self) -> None: for exp, got in zip(expected_tensors, static_tensors): assert torch.equal(got, exp) + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_group_quantize_reuses_destination_and_honors_noop(self) -> None: + """Verify the in-place and graph-safe contracts of grouped MXFP8 quantization. + + A captured training graph records the addresses of all MXFP8 weight storage, not just + the Python ``GroupedTensor`` object. Re-quantizing an updated BF16 weight must therefore + refresh the existing FP8 payloads and scales without replacing any allocation. During + graph replay, ``noop_flag`` must additionally allow the captured quantization kernel to + execute as a no-op on microbatches that should reuse the cached weight. + """ + num_tensors = 2 + # Treat two contiguous [256, 256] matrices as one grouped BF16 source. These dimensions + # satisfy the MXFP8 block-alignment requirements in both rowwise and columnwise layouts. + shape = (num_tensors * 256, 256) + quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) + + # Forward GEMM consumes rowwise weight storage, while backward dgrad consumes columnwise + # storage. Both representations and both scale tensors must be updated and kept stable. + quantizer.set_usage(rowwise=True, columnwise=True) + # Store the quantized output in GEMM-ready (swizzled) scale layout. This makes the test + # cover the exact destination format used by cached MXFP8 model weights. + quantizer.optimize_for_gemm = True + + # The first call owns allocation: it creates the persistent destination that subsequent + # optimizer updates and CUDA graph replays must modify in place. + source = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + output = tex.group_quantize(source, quantizer, num_tensors, None) + + # CUDA graphs retain these raw addresses. Object identity alone is insufficient because + # replacing one internal payload or scale allocation would leave the graph with a stale + # pointer even if ``output`` remained the same Python object. + pointers = ( + output.rowwise_data.data_ptr(), + output.columnwise_data.data_ptr(), + output.scale_inv.data_ptr(), + output.columnwise_scale_inv.data_ptr(), + ) + + # Eager in-place update: quantize new BF16 values into the previously allocated output. + # An independently allocated quantization provides the numerical reference. + updated_source = source + 1 + updated = tex.group_quantize( + updated_source, + quantizer, + num_tensors, + None, + output=output, + ) + reference = tex.group_quantize(updated_source, quantizer, num_tensors, None) + + # The API must return the caller-provided destination and preserve every captured address. + assert updated is output + assert pointers == ( + output.rowwise_data.data_ptr(), + output.columnwise_data.data_ptr(), + output.scale_inv.data_ptr(), + output.columnwise_scale_inv.data_ptr(), + ) + + # In-place quantization must still produce exactly the same FP8 bytes and scales as a + # normal allocating call, for both GEMM orientations. + assert torch.equal(output.rowwise_data, reference.rowwise_data) + assert torch.equal(output.columnwise_data, reference.columnwise_data) + assert torch.equal(output.scale_inv, reference.scale_inv) + assert torch.equal(output.columnwise_scale_inv, reference.columnwise_scale_inv) + + # Eager no-op: a nonzero device flag tells the kernel to preserve the cached destination. + # Clone all four buffers so this checks payloads and metadata independently. + old_buffers = ( + output.rowwise_data.clone(), + output.columnwise_data.clone(), + output.scale_inv.clone(), + output.columnwise_scale_inv.clone(), + ) + noop = torch.ones(1, dtype=torch.float32, device="cuda") + + # The source is deliberately different. If the kernel ignores ``noop_flag``, at least + # one payload or scale tensor below will change and expose the failure. + tex.group_quantize( + updated_source * 2, + quantizer, + num_tensors, + None, + noop_flag=noop, + output=output, + ) + assert torch.equal(output.rowwise_data, old_buffers[0]) + assert torch.equal(output.columnwise_data, old_buffers[1]) + assert torch.equal(output.scale_inv, old_buffers[2]) + assert torch.equal(output.columnwise_scale_inv, old_buffers[3]) + + # Capture one fixed launch. ``graph_source``, ``noop``, and ``output`` keep stable device + # addresses; replay changes only their contents. This is how weight caching works when + # different microbatches replay the same CUDA graph. + graph_source = updated_source.clone() + # zero means "perform quantization"; nonzero means "leave destination untouched". + noop.zero_() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + tex.group_quantize( + graph_source, + quantizer, + num_tensors, + None, + noop_flag=noop, + output=output, + ) + + # Replay with new source contents and noop=0. The captured kernel must refresh the same + # destination allocations, and the result must equal a fresh eager quantization. + graph_source.copy_(source - 1) + noop.zero_() + graph.replay() + torch.cuda.synchronize() + reference = tex.group_quantize(graph_source, quantizer, num_tensors, None) + assert torch.equal(output.rowwise_data, reference.rowwise_data) + assert torch.equal(output.columnwise_data, reference.columnwise_data) + assert torch.equal(output.scale_inv, reference.scale_inv) + assert torch.equal(output.columnwise_scale_inv, reference.columnwise_scale_inv) + + # Snapshot the successfully refreshed cache before testing the captured no-op branch. + old_buffers = ( + output.rowwise_data.clone(), + output.columnwise_data.clone(), + output.scale_inv.clone(), + output.columnwise_scale_inv.clone(), + ) + + # Replay the exact same captured graph with different source data but noop=1. The launch + # still occurs, which is required for CUDA graph structural consistency, but the kernel + # must not mutate any cached FP8 payload or scale buffer. + graph_source.copy_(source + 2) + noop.fill_(1) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(output.rowwise_data, old_buffers[0]) + assert torch.equal(output.columnwise_data, old_buffers[1]) + assert torch.equal(output.scale_inv, old_buffers[2]) + assert torch.equal(output.columnwise_scale_inv, old_buffers[3]) + @pytest.mark.parametrize("block_scaling_dim", [1, 2], ids=["1D", "2D"]) @pytest.mark.parametrize("direction", ["rowwise", "columnwise"]) @pytest.mark.skipif( diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py new file mode 100644 index 0000000000..74ec0a05ec --- /dev/null +++ b/tests/pytorch/test_hybrid_quantization.py @@ -0,0 +1,7388 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for hybrid quantization (mixed rowwise/columnwise formats).""" + +import io +import warnings +import pytest +import torch + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex + +from hybrid_quantization_utils import ( + as_data_tensor_tuple as _as_data_tensor_tuple, + assert_hybrid_tensor_exact as _assert_hybrid_tensor_exact, + assert_nested_state_exact as _assert_nested_state_exact, + assert_storage_data_exact as _assert_storage_data_exact, + fp8_e4m3_factory as _fp8_row_factory, + fp8_e5m2_factory as _fp8_grad_factory, + hybrid_block_fp8_e4m3_qfactory as _hybrid_block_fp8_qfactory, + hybrid_custom_recipe as _hybrid_custom_recipe, + hybrid_fp8_current_e5m2_grads_qfactory as _hybrid_fp8_current_qfactory, + make_fp8_quantizer as _make_fp8_quantizer, + make_hybrid_quantizer_fp8_row_fp4_col as _make_hybrid_quantizer_fp8_row_fp4_col, + make_nvfp4_quantizer as _make_nvfp4_quantizer, + make_role_aware_quantizer as _make_role_aware_quantizer, + mxfp8_e4m3_factory as _mxfp8_factory, + nvfp4_linear_mxfp8_dpa_test_factory as _nvfp4_linear_mxfp8_dpa_factory, + snapshot_storage_tensor_metadata as _snapshot_storage_tensor_metadata, +) +from transformer_engine.common import recipe +from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + nvfp4_factory, +) +from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + mxfp8_fwd_nvfp4_bwd_factory, + nvfp4_linear_fp8_dpa_factory, +) +from transformer_engine.pytorch import ( + autocast, + quantized_model_init, + Linear, + LayerNormLinear, + LayerNormMLP, + TransformerLayer, + GroupedLinear, + Float8Quantizer, + Float8CurrentScalingQuantizer, + MXFP8Quantizer, + Float8BlockQuantizer, + NVFP4Quantizer, + HybridQuantizer, + HybridQuantizedTensor, + IdentityQuantizer, + HybridQuantizedTensorStorage, + Float8Tensor, + Float8TensorStorage, + NVFP4Tensor, + NVFP4TensorStorage, + QuantizedTensor, +) +from transformer_engine.pytorch.cpp_extensions.gemm import ( + _unwrap_tensor, + _validate_native_gemm_output_quantizer, +) +from transformer_engine.pytorch.quantized_tensor import Quantizer +from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported + +_fp8_col_factory = _fp8_row_factory + +fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) +nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) +mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) +fp8_block_scaling_available, reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available( + return_reason=True +) + +_COLUMNWISE_ONLY_PER_TENSOR_FP8_ERROR = ( + "Columnwise-only per-tensor FP8 quantization is not implemented" +) + +_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 = pytest.mark.xfail( + condition=not is_non_tn_fp8_gemm_supported(), + raises=NotImplementedError, + strict=True, + reason=( + "Hopper does not yet support columnwise-only per-tensor FP8 quantization; " + "tracked by NVIDIA/TransformerEngine#3158" + ), +) + +_XFAIL_SELECTIVE_ATTENTION_RECOMPUTE = pytest.mark.xfail( + condition=is_non_tn_fp8_gemm_supported(), + raises=AttributeError, + strict=True, + reason=( + "Selective attention recompute currently consumes LayerNormLinear's saved " + "tensor objects twice; this also fails with BF16 and built-in FP8 recipes." + ), +) + +requires_fp8 = pytest.mark.skipif( + not fp8_available, + reason=f"FP8: {reason_for_no_fp8}", +) + +requires_nvfp4 = pytest.mark.skipif( + not nvfp4_available, + reason=f"NVFP4: {reason_for_no_nvfp4}", +) + +requires_fp8_and_nvfp4 = pytest.mark.skipif( + not (fp8_available and nvfp4_available), + reason=f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", +) + +requires_mxfp8_and_nvfp4 = pytest.mark.skipif( + not (mxfp8_available and nvfp4_available), + reason=f"MXFP8: {reason_for_no_mxfp8}; NVFP4: {reason_for_no_nvfp4}", +) + + +def test_native_gemm_output_quantizer_support_is_opt_in(): + """Unknown quantizers must be rejected before entering the native C++ path.""" + _validate_native_gemm_output_quantizer(None) + + unknown_quantizer = Quantizer(rowwise=True, columnwise=False) + with pytest.raises( + NotImplementedError, + match="Quantizer is not supported as a native GEMM output quantizer", + ): + _validate_native_gemm_output_quantizer(unknown_quantizer) + + +def test_hybrid_storage_snapshots_parent_quantizer(): + """Caller mutations must not change an existing Hybrid tensor's behavior.""" + quantizer = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + tensor = quantizer(torch.ones((2, 2))) + + assert tensor._quantizer is not quantizer + assert tensor._quantizer.rowwise_quantizer is not quantizer.rowwise_quantizer + assert tensor._quantizer.columnwise_quantizer is not quantizer.columnwise_quantizer + + # Make the two Identity directions independent so a skipped update is observable. + tensor._rowwise_storage._hp_data = tensor._rowwise_storage._hp_data.clone() + tensor._columnwise_storage._hp_data = tensor._columnwise_storage._hp_data.clone() + + quantizer.set_usage(rowwise=False, columnwise=True) + tensor.copy_(torch.full(tensor.shape, 2.0)) + + assert tensor._quantizer.get_usages() == {"rowwise": True, "columnwise": True} + torch.testing.assert_close( + tensor._rowwise_storage.dequantize(), + torch.full(tensor.shape, 2.0), + ) + torch.testing.assert_close( + tensor._columnwise_storage.dequantize(), + torch.full(tensor.shape, 2.0), + ) + + +def _make_hybrid_quantizer_fp4_row_fp8_col(): + """NVFP4 rowwise + FP8 columnwise (reversed direction).""" + return HybridQuantizer( + rowwise_quantizer=_make_nvfp4_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + + +def _clone_nested_state(value): + """Clone tensors in nested checkpoint state so later steps cannot mutate snapshots.""" + if isinstance(value, torch.Tensor): + return value.detach().clone() + if isinstance(value, dict): + return {key: _clone_nested_state(item) for key, item in value.items()} + if isinstance(value, list): + return [_clone_nested_state(item) for item in value] + if isinstance(value, tuple): + return tuple(_clone_nested_state(item) for item in value) + return value + + +def _snapshot_model_parameters(model): + """Capture normal values and all tensor metadata in both hybrid directions.""" + snapshot = {} + for name, param in model.named_parameters(): + if isinstance(param, HybridQuantizedTensor): + snapshot[name] = { + "rowwise": _snapshot_storage_tensor_metadata(param.rowwise_sub_storage, clone=True), + "columnwise": _snapshot_storage_tensor_metadata( + param.columnwise_sub_storage, clone=True + ), + } + else: + snapshot[name] = param.detach().clone() + return snapshot + + +class _CountingIdentityQuantizer(IdentityQuantizer): + """Replay-safe identity quantizer that counts quantize calls.""" + + def __init__(self, counter, *, dtype=None, rowwise=True, columnwise=True): + super().__init__(dtype=dtype, rowwise=rowwise, columnwise=columnwise) + self.counter = counter + + def copy(self): + quantizer = type(self)( + self.counter, + dtype=self.dtype, + rowwise=self.rowwise_usage, + columnwise=self.columnwise_usage, + ) + quantizer.internal = self.internal + quantizer.optimize_for_gemm = self.optimize_for_gemm + return quantizer + + def quantize_impl(self, tensor): + self.counter["calls"] += 1 + return super().quantize_impl(tensor) + + +class _CountingUnsafeIdentityQuantizer(_CountingIdentityQuantizer): + """Non-replay-safe identity quantizer that counts quantize calls.""" + + def is_requantization_safe(self): + self.counter["safety_calls"] = self.counter.get("safety_calls", 0) + 1 + return False + + +class _CountingPythonQuantizer(Quantizer): + """Custom Python quantizer used to verify fallback dispatch.""" + + def __init__(self, calls, *, rowwise=True, columnwise=True): + super().__init__(rowwise=rowwise, columnwise=columnwise) + self.calls = calls + + def copy(self): + quantizer = type(self)( + self.calls, + rowwise=self.rowwise_usage, + columnwise=self.columnwise_usage, + ) + quantizer.internal = self.internal + quantizer.optimize_for_gemm = self.optimize_for_gemm + return quantizer + + def quantize_impl(self, tensor): + self.calls.append(tensor) + fallback = IdentityQuantizer( + rowwise=self.rowwise_usage, + columnwise=self.columnwise_usage, + ) + fallback.internal = self.internal + return fallback.quantize_impl(tensor) + + +@requires_mxfp8_and_nvfp4 +class TestComposerStyleFactory: + """Composer 2-style row-scaled NVFP4 forward + MXFP8 backward dispatch.""" + + @staticmethod + def _factory(role): + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + nvfp4_row_scaled_fwd_mxfp8_bwd_factory, + ) + + return nvfp4_row_scaled_fwd_mxfp8_bwd_factory(role) + + @pytest.mark.parametrize( + "tensor_type,row_scaled", + [("input", True), ("weight", False)], + ) + def test_grouped_linear_forward_roles_use_nvfp4_rowwise_mxfp8_columnwise( + self, tensor_type, row_scaled + ): + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = self._factory( + QuantizerRole(module_type="grouped_linear", tensor_type=tensor_type) + ) + + assert isinstance(quantizer, HybridQuantizer) + assert quantizer.columnwise_source == "rowwise_dequantized" + assert isinstance(quantizer.rowwise_quantizer, NVFP4Quantizer) + assert isinstance(quantizer.columnwise_quantizer, MXFP8Quantizer) + assert quantizer.rowwise_quantizer.row_scaled_nvfp4 is row_scaled + assert quantizer.rowwise_quantizer.with_rht is False + assert quantizer.rowwise_quantizer.with_post_rht_amax is False + assert quantizer.rowwise_quantizer.with_2d_quantization is False + assert quantizer.rowwise_quantizer.stochastic_rounding is False + assert quantizer.rowwise_quantizer.rowwise_usage is True + assert quantizer.rowwise_quantizer.columnwise_usage is False + assert quantizer.columnwise_quantizer.rowwise_usage is False + assert quantizer.columnwise_quantizer.columnwise_usage is True + + @pytest.mark.parametrize("tensor_type", ["input", "output", "weight"]) + def test_regular_linear_forward_roles_fall_back_to_mxfp8(self, tensor_type): + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = self._factory(QuantizerRole(module_type="linear", tensor_type=tensor_type)) + + assert isinstance(quantizer, MXFP8Quantizer) + + @pytest.mark.parametrize("module_type", ["linear", "grouped_linear"]) + @pytest.mark.parametrize("tensor_type", ["grad_output", "grad_input"]) + def test_backward_roles_use_mxfp8(self, module_type, tensor_type): + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = self._factory(QuantizerRole(module_type=module_type, tensor_type=tensor_type)) + + assert isinstance(quantizer, MXFP8Quantizer) + + def test_non_linear_roles_fall_back_to_mxfp8(self): + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = self._factory(QuantizerRole(module_type="dpa", tensor_type="qkv")) + + assert isinstance(quantizer, MXFP8Quantizer) + + +@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") +class TestMXFP8FwdHighPrecisionBwdFactory: + """MXFP8 forward + dequantized backward qfactory dispatch.""" + + @staticmethod + def _factory(role): + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + mxfp8_fwd_high_precision_bwd_factory, + ) + + return mxfp8_fwd_high_precision_bwd_factory(role) + + @pytest.mark.parametrize("module_type", ["linear", "grouped_linear"]) + @pytest.mark.parametrize("tensor_type", ["input", "weight"]) + def test_forward_roles_are_requantization_safe(self, module_type, tensor_type): + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = self._factory(QuantizerRole(module_type=module_type, tensor_type=tensor_type)) + + assert isinstance(quantizer, HybridQuantizer) + assert quantizer.columnwise_source == "rowwise_dequantized" + assert isinstance(quantizer.rowwise_quantizer, MXFP8Quantizer) + assert isinstance(quantizer.columnwise_quantizer, IdentityQuantizer) + assert quantizer.is_requantization_safe() is True + + @pytest.mark.parametrize("module_type", ["linear", "grouped_linear"]) + def test_grad_output_role_is_requantization_safe(self, module_type): + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = self._factory(QuantizerRole(module_type=module_type, tensor_type="grad_output")) + + assert isinstance(quantizer, IdentityQuantizer) + assert quantizer.is_requantization_safe() is True + + +@requires_nvfp4 +class TestNVFP41DWeightFactory: + """Base NVFP4 recipe with W.T sourced from dequantized 1D W.""" + + @staticmethod + def _factory(role): + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + nvfp4_1d_weight_factory, + ) + + return nvfp4_1d_weight_factory(role) + + @staticmethod + def _assert_plain_1d_nvfp4(quantizer): + assert isinstance(quantizer, NVFP4Quantizer) + assert quantizer.row_scaled_nvfp4 is False + assert quantizer.with_rht is False + assert quantizer.with_post_rht_amax is False + assert quantizer.with_2d_quantization is False + assert quantizer.stochastic_rounding is False + + @pytest.mark.parametrize("module_type", ["linear", "grouped_linear"]) + def test_weight_roles_use_rowwise_dequantized_source(self, module_type): + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = self._factory(QuantizerRole(module_type=module_type, tensor_type="weight")) + + assert isinstance(quantizer, HybridQuantizer) + assert quantizer.columnwise_source == "rowwise_dequantized" + self._assert_plain_1d_nvfp4(quantizer.rowwise_quantizer) + self._assert_plain_1d_nvfp4(quantizer.columnwise_quantizer) + assert quantizer.rowwise_quantizer.rowwise_usage is True + assert quantizer.rowwise_quantizer.columnwise_usage is False + assert quantizer.columnwise_quantizer.rowwise_usage is False + assert quantizer.columnwise_quantizer.columnwise_usage is True + + @staticmethod + def _assert_matches_base_nvfp4_recipe(quantizer, expected): + assert isinstance(quantizer, NVFP4Quantizer) + assert isinstance(expected, NVFP4Quantizer) + assert quantizer.dtype == expected.dtype + assert quantizer.rowwise_usage == expected.rowwise_usage + assert quantizer.columnwise_usage == expected.columnwise_usage + assert quantizer.with_rht == expected.with_rht + assert quantizer.with_post_rht_amax == expected.with_post_rht_amax + assert quantizer.with_2d_quantization == expected.with_2d_quantization + assert quantizer.stochastic_rounding == expected.stochastic_rounding + assert quantizer.row_scaled_nvfp4 == expected.row_scaled_nvfp4 + + @pytest.mark.parametrize("module_type", ["linear", "grouped_linear"]) + @pytest.mark.parametrize("tensor_type", ["input", "output", "grad_output", "grad_input"]) + def test_non_weight_linear_roles_match_base_nvfp4_recipe(self, module_type, tensor_type): + from transformer_engine.pytorch.quantization import QuantizerRole + from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + nvfp4_factory, + ) + + role = QuantizerRole(module_type=module_type, tensor_type=tensor_type) + quantizer = self._factory(role) + expected = nvfp4_factory(role) + + self._assert_matches_base_nvfp4_recipe(quantizer, expected) + + def test_non_linear_roles_match_base_nvfp4_recipe(self): + from transformer_engine.pytorch.quantization import QuantizerRole + from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + nvfp4_factory, + ) + + role = QuantizerRole(module_type="dpa", tensor_type="qkv") + quantizer = self._factory(role) + expected = nvfp4_factory(role) + + self._assert_matches_base_nvfp4_recipe(quantizer, expected) + + def test_weight_quantizer_produces_both_nvfp4_storages(self): + from transformer_engine.pytorch.quantization import QuantizerRole + + torch.manual_seed(2026) + src = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + quantizer = self._factory(QuantizerRole(module_type="linear", tensor_type="weight")) + + out = quantizer.quantize(src) + + assert isinstance(out.rowwise_sub_storage, (NVFP4TensorStorage, NVFP4Tensor)) + assert isinstance(out.columnwise_sub_storage, (NVFP4TensorStorage, NVFP4Tensor)) + + +@requires_fp8_and_nvfp4 +class TestHybridQuantizerConstruction: + """Test construction and basic properties of hybrid quantizer.""" + + def test_creation(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + assert isinstance(hq, HybridQuantizer) + assert hq.rowwise_usage is True + assert hq.columnwise_usage is True + assert isinstance(hq.rowwise_quantizer, Float8CurrentScalingQuantizer) + assert isinstance(hq.columnwise_quantizer, NVFP4Quantizer) + + def test_hybrid_storage_and_tensor_require_parent_quantizer(self): + with pytest.raises(TypeError, match="requires a parent HybridQuantizer"): + HybridQuantizedTensorStorage( + rowwise_storage=None, + columnwise_storage=None, + quantizer=None, + ) + with pytest.raises(TypeError, match="requires a parent HybridQuantizer"): + HybridQuantizedTensor( + shape=(1, 1), + dtype=torch.bfloat16, + rowwise_storage=None, + columnwise_storage=None, + quantizer=None, + ) + + def test_rejects_same_sub_quantizer_instance_for_both_directions(self): + quantizer = _make_fp8_quantizer() + + with pytest.raises(ValueError, match="requires distinct rowwise and columnwise"): + HybridQuantizer(rowwise_quantizer=quantizer, columnwise_quantizer=quantizer) + + assert quantizer.rowwise_usage is True + assert quantizer.columnwise_usage is True + + def test_compatible_recipe_is_custom_recipe(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + assert hq._get_compatible_recipe() is recipe.CustomRecipe + + def test_supports_only_rowwise_all_gather_nvfp4_columnwise(self): + """NVFP4 columnwise sub-quantizer forces rowwise-only AG. + + ``NVFP4Tensor.dequantize()`` raises ``NotImplementedError`` for + columnwise-only data, so the BF16 fallback in + ``gather_along_first_dim`` cannot operate on a columnwise-only + NVFP4 hybrid sub-storage. ``HybridQuantizer.supports_only_rowwise_all_gather`` + must return True in this case so ``_linear_forward_impl`` / + ``_linear_backward`` preserve rowwise data (which NVFP4 can + dequantize) instead. + """ + hq = HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + assert hq.supports_only_rowwise_all_gather() is True + + def test_supports_only_rowwise_all_gather_mxfp8_both(self): + """MXFP8 in both directions → columnwise dequant works → default + False so the save-columnwise (for wgrad) path stays active.""" + if not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") + hq = HybridQuantizer( + rowwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + columnwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + ) + assert hq.supports_only_rowwise_all_gather() is False + + def test_supports_only_rowwise_all_gather_fp8_current_propagates(self): + """Float8CurrentScalingQuantizer returns True for its own flag; + hybrid must propagate (not swallow) that semantics.""" + hq = HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + assert hq.supports_only_rowwise_all_gather() is True + + def test_supports_only_rowwise_all_gather_nvfp4_both(self): + """NVFP4 in both directions → columnwise sub-quantizer is NVFP4 + → forces rowwise-only AG regardless of rowwise flag.""" + hq = HybridQuantizer( + rowwise_quantizer=_make_nvfp4_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + assert hq.supports_only_rowwise_all_gather() is True + + +@requires_fp8 +class TestHybridColumnwiseSource: + """Test columnwise source provenance in HybridQuantizer.""" + + @pytest.fixture + def input_tensor(self): + torch.manual_seed(90210) + return torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + + @staticmethod + def _make_quantizer(columnwise_source="original"): + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source=columnwise_source, + ) + + def test_default_columnwise_source_original(self, input_tensor): + hq = self._make_quantizer() + out = hq.quantize(input_tensor) + + assert hq.columnwise_source == "original" + torch.testing.assert_close( + out.columnwise_sub_storage.dequantize(), input_tensor, rtol=0.0, atol=0.0 + ) + + def test_invalid_columnwise_source_raises(self): + with pytest.raises(ValueError, match="columnwise_source"): + HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="dequantized", + ) + + def test_default_quantizer_is_requantization_safe(self): + assert IdentityQuantizer().is_requantization_safe() is True + + @pytest.mark.parametrize("columnwise_source", ("original", "rowwise_dequantized")) + def test_hybrid_requantization_safe_with_safe_sub_quantizers(self, columnwise_source): + hq = self._make_quantizer(columnwise_source=columnwise_source) + + assert hq.is_requantization_safe() is True + + def test_hybrid_requantization_checks_requested_sub_quantizers(self): + from transformer_engine.pytorch.module._common import ( + can_reconstruct_wgrad_input_from_original, + ) + + hq = HybridQuantizer( + rowwise_quantizer=_CountingUnsafeIdentityQuantizer({"calls": 0}), + columnwise_quantizer=_CountingUnsafeIdentityQuantizer({"calls": 0}), + columnwise_source="original", + ) + + assert hq.is_requantization_safe() is False + assert can_reconstruct_wgrad_input_from_original(hq) is True + + @pytest.mark.parametrize("columnwise_source", ("original", "rowwise_dequantized")) + @pytest.mark.parametrize("rowwise_safe", (True, False)) + @pytest.mark.parametrize("columnwise_safe", (True, False)) + def test_requantization_and_wgrad_reconstruction_policy_matrix( + self, + columnwise_source, + rowwise_safe, + columnwise_safe, + ): + from transformer_engine.pytorch.module._common import ( + can_reconstruct_wgrad_input_from_original, + ) + + rowwise_cls = ( + _CountingIdentityQuantizer if rowwise_safe else _CountingUnsafeIdentityQuantizer + ) + columnwise_cls = ( + _CountingIdentityQuantizer if columnwise_safe else _CountingUnsafeIdentityQuantizer + ) + hq = HybridQuantizer( + rowwise_quantizer=rowwise_cls({"calls": 0}), + columnwise_quantizer=columnwise_cls({"calls": 0}), + columnwise_source=columnwise_source, + ) + + # General requantization reproduces both requested representations. + assert hq.is_requantization_safe() is (rowwise_safe and columnwise_safe) + # Wgrad reconstruction only repeats the rowwise stage for a + # rowwise-dequantized columnwise source. + expected_reconstruction = columnwise_source == "original" or rowwise_safe + assert can_reconstruct_wgrad_input_from_original(hq) is expected_reconstruction + + def test_hybrid_rowwise_source_requires_safe_rowwise_quantizer(self): + hq = HybridQuantizer( + rowwise_quantizer=_CountingUnsafeIdentityQuantizer({"calls": 0}), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + + assert hq.is_requantization_safe() is False + + def test_copy_preserves_columnwise_source(self): + hq = self._make_quantizer(columnwise_source="rowwise_dequantized") + hq.set_usage(rowwise=False, columnwise=True) + + copied = hq.copy() + + assert copied.columnwise_source == "rowwise_dequantized" + assert copied.rowwise_usage is False + assert copied.columnwise_usage is True + + def test_rowwise_dequantized_identity_columnwise_matches_rowwise(self, input_tensor): + hq = self._make_quantizer(columnwise_source="rowwise_dequantized") + out = hq.quantize(input_tensor) + + rowwise_dq = out.rowwise_sub_storage.dequantize() + columnwise_dq = out.columnwise_sub_storage.dequantize() + torch.testing.assert_close(columnwise_dq, rowwise_dq, rtol=0.0, atol=0.0) + + def test_internal_rowwise_storage_preserves_input_dtype(self, input_tensor): + hq = self._make_quantizer(columnwise_source="rowwise_dequantized") + hq.rowwise_quantizer.internal = True + + columnwise_src = hq._columnwise_src_from_rowwise(input_tensor, None) + + assert columnwise_src.dtype == input_tensor.dtype + + def test_columnwise_only_uses_transient_rowwise_source(self, input_tensor): + hq = self._make_quantizer(columnwise_source="rowwise_dequantized") + hq.set_usage(rowwise=False, columnwise=True) + expected = _make_fp8_quantizer().quantize(input_tensor).dequantize() + + out = hq.quantize(input_tensor) + + assert out.rowwise_sub_storage is None + assert out.columnwise_sub_storage is not None + torch.testing.assert_close( + out.columnwise_sub_storage.dequantize(), expected, rtol=0.0, atol=0.0 + ) + + def test_update_quantized_columnwise_only_uses_transient_rowwise_source(self, input_tensor): + hq = self._make_quantizer(columnwise_source="rowwise_dequantized") + dst = hq.quantize(input_tensor) + rowwise_before = dst.rowwise_sub_storage.dequantize().clone() + new_src = torch.randn_like(input_tensor) * 8 + expected = _make_fp8_quantizer().quantize(new_src).dequantize() + + hq.set_usage(rowwise=False, columnwise=True) + hq.update_quantized(new_src, dst) + + torch.testing.assert_close( + dst.rowwise_sub_storage.dequantize(), rowwise_before, rtol=0.0, atol=0.0 + ) + torch.testing.assert_close( + dst.columnwise_sub_storage.dequantize(), expected, rtol=0.0, atol=0.0 + ) + + def test_update_quantized_uses_updated_rowwise_storage(self, input_tensor): + hq = self._make_quantizer(columnwise_source="rowwise_dequantized") + dst = hq.quantize(input_tensor) + new_src = torch.randn_like(input_tensor) * 8 + + hq.set_usage(rowwise=True, columnwise=True) + hq.update_quantized(new_src, dst) + + torch.testing.assert_close( + dst.columnwise_sub_storage.dequantize(), + dst.rowwise_sub_storage.dequantize(), + rtol=0.0, + atol=0.0, + ) + + +@requires_fp8 +class TestHybridSaveOriginalInputPolicy: + """Module-level save_original_input policy for hybrid qfactory inputs.""" + + def test_linear_high_precision_override_skips_requantization_safety_veto(self): + counters = {"rowwise": {"calls": 0}, "columnwise": {"calls": 0}} + custom_recipe = self._counting_identity_hybrid_recipe(counters) + custom_recipe.backward_override = "high_precision" + model = Linear( + 128, + 128, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=False, + ).cuda() + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + with autocast(enabled=True, recipe=custom_recipe): + out = model(inp) + out.float().sum().backward() + + assert counters["rowwise"].get("safety_calls", 0) == 0 + + def test_linear_custom_delayed_scaling_disables_save_original_input(self): + from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + delayed_scaling_factory, + ) + + model = Linear( + 128, + 128, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=True, + ).cuda() + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + custom_recipe = recipe.CustomRecipe(qfactory=delayed_scaling_factory) + + with pytest.warns( + UserWarning, + match="save_original_input is incompatible with delayed-scaling quantizers", + ): + with autocast(enabled=True, recipe=custom_recipe): + out = model(inp) + out.float().sum().backward() + + def test_linear_builtin_delayed_scaling_rejects_save_original_input(self): + model = Linear( + 128, + 128, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=True, + ).cuda() + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + with pytest.raises( + ValueError, + match="DelayedScaling recipe is not supported with save_original_input", + ): + with autocast(enabled=True, recipe=recipe.DelayedScaling()): + model(inp) + + def test_grouped_linear_classifies_requantization_safety_once_per_generation(self): + counters = [{"calls": 0}, {"calls": 0}] + input_quantizers = [_CountingUnsafeIdentityQuantizer(counter) for counter in counters] + generation = [] + for input_quantizer in input_quantizers: + generation.extend((input_quantizer, IdentityQuantizer(), IdentityQuantizer())) + + module = GroupedLinear( + 2, + 16, + 16, + bias=False, + device="meta", + ) + module.quantizers["scaling_fwd"] = generation + + module._validate_quantizer_generation(fwd=True) + assert module._unsafe_requantization_input_quantizer is input_quantizers[0] + assert [counter.get("safety_calls", 0) for counter in counters] == [1, 0] + + # The generation list is stable between forwards, so the O(1) identity + # guard must avoid re-running capability checks. + module._validate_quantizer_generation(fwd=True) + assert [counter.get("safety_calls", 0) for counter in counters] == [1, 0] + + @staticmethod + def _counting_identity_hybrid_recipe( + counters, + *, + columnwise_source="rowwise_dequantized", + rowwise_safe=False, + columnwise_safe=True, + ): + def factory(role): + if role is not None and role.module_type == "linear" and role.tensor_type == "input": + rowwise_cls = ( + _CountingIdentityQuantizer if rowwise_safe else _CountingUnsafeIdentityQuantizer + ) + columnwise_cls = ( + _CountingIdentityQuantizer + if columnwise_safe + else _CountingUnsafeIdentityQuantizer + ) + return HybridQuantizer( + rowwise_quantizer=rowwise_cls(counters["rowwise"]), + columnwise_quantizer=columnwise_cls(counters["columnwise"]), + columnwise_source=columnwise_source, + ) + return IdentityQuantizer() + + return recipe.CustomRecipe(qfactory=factory) + + def test_linear_save_original_input_veto_uses_saved_forward_quantized_input(self): + torch.manual_seed(205) + counters = {"rowwise": {"calls": 0}, "columnwise": {"calls": 0}} + model = Linear( + 128, + 128, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=True, + ).cuda() + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + with pytest.warns(UserWarning, match="Ignoring save_original_input=True"): + with autocast( + enabled=True, + recipe=self._counting_identity_hybrid_recipe(counters), + ): + out = model(inp) + + calls_after_forward = counters["rowwise"]["calls"] + assert calls_after_forward > 0 + + out.float().sum().backward() + + assert counters["rowwise"]["calls"] == calls_after_forward + + @pytest.mark.parametrize( + "columnwise_source,rowwise_safe,columnwise_safe,optimization_enabled," + "expected_rowwise_calls", + ( + ("original", True, True, True, 1), + ("original", False, False, True, 1), + ("rowwise_dequantized", True, True, True, 2), + ("rowwise_dequantized", True, False, True, 2), + ("rowwise_dequantized", False, True, False, 1), + ), + ) + def test_linear_save_original_input_policy_is_bitwise_exact( + self, + columnwise_source, + rowwise_safe, + columnwise_safe, + optimization_enabled, + expected_rowwise_calls, + ): + torch.manual_seed(206) + in_features, out_features, batch = 128, 128, 32 + reference = Linear( + in_features, + out_features, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=False, + ).cuda() + candidate = Linear( + in_features, + out_features, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=True, + ).cuda() + candidate.load_state_dict(reference.state_dict()) + + base_input = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + reference_input = base_input.clone().detach().requires_grad_(True) + candidate_input = base_input.clone().detach().requires_grad_(True) + reference_counters = { + "rowwise": {"calls": 0}, + "columnwise": {"calls": 0}, + } + candidate_counters = { + "rowwise": {"calls": 0}, + "columnwise": {"calls": 0}, + } + recipe_kwargs = { + "columnwise_source": columnwise_source, + "rowwise_safe": rowwise_safe, + "columnwise_safe": columnwise_safe, + } + + with autocast( + enabled=True, + recipe=self._counting_identity_hybrid_recipe(reference_counters, **recipe_kwargs), + ): + reference_output = reference(reference_input) + + candidate_recipe = self._counting_identity_hybrid_recipe( + candidate_counters, + **recipe_kwargs, + ) + if optimization_enabled: + with autocast(enabled=True, recipe=candidate_recipe): + candidate_output = candidate(candidate_input) + else: + with pytest.warns(UserWarning, match="Ignoring save_original_input=True"): + with autocast(enabled=True, recipe=candidate_recipe): + candidate_output = candidate(candidate_input) + + assert torch.equal(reference_output, candidate_output) + + reference_output.float().sum().backward() + candidate_output.float().sum().backward() + + assert reference_input.grad is not None and candidate_input.grad is not None + assert torch.equal(reference_input.grad, candidate_input.grad) + reference_weight_grad = dict(reference.named_parameters())["weight"].grad + candidate_weight_grad = dict(candidate.named_parameters())["weight"].grad + assert reference_weight_grad is not None and candidate_weight_grad is not None + assert torch.equal(reference_weight_grad, candidate_weight_grad) + + # The columnwise quantizer is always invoked exactly once. Only a + # rowwise-dequantized reconstruction safely repeats the rowwise stage. + assert candidate_counters["rowwise"]["calls"] == expected_rowwise_calls + assert candidate_counters["columnwise"]["calls"] == 1 + + +@requires_fp8_and_nvfp4 +class TestHybridQuantize: + """Test quantization via HybridQuantizer.""" + + @pytest.fixture + def input_tensor(self): + torch.manual_seed(42) + return torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + + def test_quantize_returns_hybrid_tensor(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + assert isinstance(result, HybridQuantizedTensor) + + def test_quantize_shape_preserved(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + assert result.shape == input_tensor.shape + + def test_quantize_dtype_preserved(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + assert result.dtype == input_tensor.dtype + + def test_sub_storage_types_fp8_row_fp4_col(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + row_storage = result.rowwise_sub_storage + col_storage = result.columnwise_sub_storage + assert isinstance(row_storage, (Float8TensorStorage, Float8Tensor)) + assert isinstance(col_storage, (NVFP4TensorStorage, NVFP4Tensor)) + + def test_sub_storage_types_reversed(self, input_tensor): + hq = _make_hybrid_quantizer_fp4_row_fp8_col() + result = hq.quantize(input_tensor) + row_storage = result.rowwise_sub_storage + col_storage = result.columnwise_sub_storage + assert isinstance(row_storage, (NVFP4TensorStorage, NVFP4Tensor)) + assert isinstance(col_storage, (Float8TensorStorage, Float8Tensor)) + + def test_quantize_internal_returns_storage(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.internal = True + result = hq.quantize(input_tensor) + assert isinstance(result, HybridQuantizedTensorStorage) + assert not isinstance(result, HybridQuantizedTensor) + hq.internal = False + + +@requires_fp8_and_nvfp4 +class TestHybridDequantize: + """Test dequantization round-trip.""" + + @pytest.fixture + def input_tensor(self): + torch.manual_seed(42) + return torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + + def test_dequantize_shape(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + dequantized = result.dequantize() + assert dequantized.shape == input_tensor.shape + + def test_dequantize_dtype(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + dequantized = result.dequantize() + assert dequantized.dtype == input_tensor.dtype + + def test_dequantize_explicit_dtype(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + dequantized = result.dequantize(dtype=torch.float32) + assert dequantized.dtype == torch.float32 + assert dequantized.shape == input_tensor.shape + + def test_dequantize_close_to_original(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + dequantized = result.dequantize() + torch.testing.assert_close( + dequantized.float(), input_tensor.float(), rtol=0.125, atol=0.0675 + ) + + def test_dequantize_reversed_close_to_original(self, input_tensor): + hq = _make_hybrid_quantizer_fp4_row_fp8_col() + result = hq.quantize(input_tensor) + dequantized = result.dequantize() + torch.testing.assert_close(dequantized.float(), input_tensor.float(), rtol=0.5, atol=1.0) + + def test_storage_dequantize(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.internal = True + result = hq.quantize(input_tensor) + dequantized = result.dequantize(dtype=torch.bfloat16) + assert dequantized.shape == input_tensor.shape + hq.internal = False + + +class TestHybridUpdateUsage: + """Test update_usage semantics and sub-storage cleanup.""" + + @pytest.fixture + def hybrid_tensor(self): + inp = torch.randn(4, 8) + hq = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + return hq.quantize(inp) + + def test_initial_usages(self, hybrid_tensor): + usages = hybrid_tensor.get_usages() + assert usages["rowwise"] is True + assert usages["columnwise"] is True + + def test_drop_rowwise(self, hybrid_tensor): + hybrid_tensor.update_usage(rowwise_usage=False) + assert hybrid_tensor.rowwise_sub_storage is None + assert hybrid_tensor.columnwise_sub_storage is not None + usages = hybrid_tensor.get_usages() + assert usages["rowwise"] is False + assert usages["columnwise"] is True + + def test_drop_columnwise(self, hybrid_tensor): + hybrid_tensor.update_usage(columnwise_usage=False) + assert hybrid_tensor.columnwise_sub_storage is None + assert hybrid_tensor.rowwise_sub_storage is not None + usages = hybrid_tensor.get_usages() + assert usages["rowwise"] is True + assert usages["columnwise"] is False + + def test_drop_both(self, hybrid_tensor): + hybrid_tensor.update_usage(rowwise_usage=False, columnwise_usage=False) + usages = hybrid_tensor.get_usages() + assert usages["rowwise"] is False + assert usages["columnwise"] is False + + def test_request_true_is_noop(self, hybrid_tensor): + row_before = hybrid_tensor.rowwise_sub_storage + col_before = hybrid_tensor.columnwise_sub_storage + hybrid_tensor.update_usage(rowwise_usage=True, columnwise_usage=True) + assert hybrid_tensor.rowwise_sub_storage is row_before + assert hybrid_tensor.columnwise_sub_storage is col_before + + def test_request_missing_columnwise_raises(self, hybrid_tensor): + hybrid_tensor.update_usage(columnwise_usage=False) + + with pytest.raises(RuntimeError, match="no columnwise sub-storage"): + hybrid_tensor.update_usage(columnwise_usage=True) + + assert hybrid_tensor.rowwise_sub_storage is not None + assert hybrid_tensor.columnwise_sub_storage is None + + def test_request_missing_rowwise_raises(self, hybrid_tensor): + hybrid_tensor.update_usage(rowwise_usage=False) + + with pytest.raises(RuntimeError, match="no rowwise sub-storage"): + hybrid_tensor.update_usage(rowwise_usage=True) + + assert hybrid_tensor.rowwise_sub_storage is None + assert hybrid_tensor.columnwise_sub_storage is not None + + def test_missing_direction_request_is_atomic(self, hybrid_tensor): + hybrid_tensor.update_usage(columnwise_usage=False) + row_before = hybrid_tensor.rowwise_sub_storage + + with pytest.raises(RuntimeError, match="no columnwise sub-storage"): + hybrid_tensor.update_usage(rowwise_usage=False, columnwise_usage=True) + + assert hybrid_tensor.rowwise_sub_storage is row_before + assert hybrid_tensor.columnwise_sub_storage is None + + def test_none_preserves_missing_direction(self, hybrid_tensor): + hybrid_tensor.update_usage(columnwise_usage=False) + row_before = hybrid_tensor.rowwise_sub_storage + + hybrid_tensor.update_usage(rowwise_usage=None, columnwise_usage=None) + + assert hybrid_tensor.rowwise_sub_storage is row_before + assert hybrid_tensor.columnwise_sub_storage is None + + def test_repr_after_drop(self, hybrid_tensor): + hybrid_tensor.update_usage(rowwise_usage=False) + r = repr(hybrid_tensor) + assert "HybridQuantizedTensor" in r + assert "rowwise=None" in r + + hybrid_tensor.update_usage(columnwise_usage=False) + r = repr(hybrid_tensor) + assert "rowwise=None" in r + assert "columnwise=None" in r + + +requires_mxfp8 = pytest.mark.skipif( + not mxfp8_available, + reason=f"MXFP8: {reason_for_no_mxfp8}", +) + + +@requires_fp8_and_nvfp4 +class TestHybridClear: + """Test HybridQuantizedTensorStorage.clear() — buffer deallocation. + + ``clear()`` is invoked by cpu_offload_v1 after the offloader has taken + its own reference to the extracted buffers, to free the GPU-resident + originals. HybridQuantizedTensorStorage delegates to each sub-storage's + own clear(), which replaces primary data buffers with empty tensors. + """ + + @pytest.fixture + def input_tensor(self): + torch.manual_seed(42) + return torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + + @staticmethod + def _primary_data_numels(sub_storage): + """Collect numel() of primary data buffers on a sub-storage. + + After ``clear()`` every entry should be 0 (native sub-storages set + ``t.data = _empty_tensor()`` on the primary buffers). + """ + if sub_storage is None: + return [] + data = sub_storage.get_data_tensors() + if not isinstance(data, tuple): + data = (data,) + return [t.numel() for t in data if t is not None] + + def test_clear_delegates_to_both_sub_storages(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + + row_before = self._primary_data_numels(ht.rowwise_sub_storage) + col_before = self._primary_data_numels(ht.columnwise_sub_storage) + assert row_before and all(n > 0 for n in row_before) + assert col_before and all(n > 0 for n in col_before) + + ht.clear() + + row_after = self._primary_data_numels(ht.rowwise_sub_storage) + col_after = self._primary_data_numels(ht.columnwise_sub_storage) + assert all(n == 0 for n in row_after) + assert all(n == 0 for n in col_after) + + @requires_mxfp8 + def test_clear_delegates_mxfp8_nvfp4(self, input_tensor): + """Per-block sub-storage path (MXFP8 rowwise + NVFP4 columnwise).""" + hq = HybridQuantizer( + rowwise_quantizer=_make_mxfp8_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + ht = hq.quantize(input_tensor) + ht.clear() + for sub in (ht.rowwise_sub_storage, ht.columnwise_sub_storage): + for n in self._primary_data_numels(sub): + assert n == 0 + + def test_clear_with_rowwise_only(self, input_tensor): + """columnwise sub-storage is None — clear() must not crash.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + ht.update_usage(columnwise_usage=False) + assert ht.columnwise_sub_storage is None + + ht.clear() + + assert all(n == 0 for n in self._primary_data_numels(ht.rowwise_sub_storage)) + + def test_clear_with_columnwise_only(self, input_tensor): + """rowwise sub-storage is None — clear() must not crash.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + ht.update_usage(rowwise_usage=False) + assert ht.rowwise_sub_storage is None + + ht.clear() + + assert all(n == 0 for n in self._primary_data_numels(ht.columnwise_sub_storage)) + + def test_clear_with_both_sub_storages_dropped(self, input_tensor): + """Both sub-storages are None — clear() must not crash.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + ht.update_usage(rowwise_usage=False, columnwise_usage=False) + assert ht.rowwise_sub_storage is None + assert ht.columnwise_sub_storage is None + + ht.clear() # must not raise + + def test_clear_is_idempotent(self, input_tensor): + """Calling clear() twice must not raise and leaves buffers empty.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + ht.clear() + ht.clear() + for sub in (ht.rowwise_sub_storage, ht.columnwise_sub_storage): + for n in self._primary_data_numels(sub): + assert n == 0 + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridTensorShapeOps: + """Shape ops that preserve supported Hybrid sub-storages.""" + + def test_fp8_current_non_noop_slice_and_narrow_preserve_hybrid(self): + torch.manual_seed(42) + x = torch.randn(64, 32, dtype=torch.bfloat16, device="cuda") + quantizer = HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + tensor = quantizer.quantize(x) + dequantized = tensor.dequantize() + + sliced = torch.ops.aten.slice.Tensor(tensor, 0, 0, 32, 1) + narrowed = tensor.narrow(0, 16, 32) + + assert isinstance(sliced, HybridQuantizedTensor) + assert isinstance(narrowed, HybridQuantizedTensor) + torch.testing.assert_close(sliced.dequantize(), dequantized[:32], rtol=0, atol=0) + torch.testing.assert_close(narrowed.dequantize(), dequantized[16:48], rtol=0, atol=0) + + def test_full_span_step_slice_is_not_treated_as_noop(self): + torch.manual_seed(42) + x = torch.randn(64, 32, dtype=torch.bfloat16, device="cuda") + quantizer = HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + tensor = quantizer.quantize(x) + dequantized = tensor.dequantize() + + sliced = torch.ops.aten.slice.Tensor(tensor, 0, 0, tensor.size(0), 2) + + assert isinstance(sliced, HybridQuantizedTensor) + torch.testing.assert_close(sliced.dequantize(), dequantized[::2], rtol=0, atol=0) + + def test_same_shape_as_strided_with_offset_is_not_treated_as_noop(self): + torch.manual_seed(42) + x = torch.randn(64, 32, dtype=torch.bfloat16, device="cuda") + quantizer = HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + tensor = quantizer.quantize(x) + dequantized = tensor.dequantize() + + base_view = torch.ops.aten.slice.Tensor(tensor, 0, 0, 63, 1) + shifted = torch.ops.aten.as_strided.default( + base_view, base_view.shape, base_view.stride(), x.stride(0) + ) + + assert isinstance(shifted, HybridQuantizedTensor) + torch.testing.assert_close(shifted.dequantize(), dequantized[1:], rtol=0, atol=0) + + +@requires_fp8_and_nvfp4 +class TestHybridDetachIsolation: + """``HybridQuantizedTensor.detach()`` must produce a hybrid whose + sub-storage wrappers are NOT shared with ``self`` — so that + ``detached.prepare_for_saving()`` does not null out fields on the + original. + + This is the property cpu_offload_v2 relies on at + ``cpu_offload.py:378-382``: + + tensor_copy = tensor.detach() + saved_tensors, _ = tensor_copy.prepare_for_saving() + """ + + @pytest.fixture + def input_tensor(self): + torch.manual_seed(42) + return torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + + def test_detach_produces_distinct_sub_storage_wrappers(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + detached = ht.detach() + + assert detached is not ht + assert detached._rowwise_storage is not ht._rowwise_storage + assert detached._columnwise_storage is not ht._columnwise_storage + + def test_detach_prepare_for_saving_does_not_affect_original(self, input_tensor): + """prepare_for_saving on the detach() copy must not null the original. + + This is the specific invariant the cpu_offload_v2 push/reload cycle + depends on. Without it, a second push on the same tensor — or even + a bare ``.device`` read during offload eligibility checks — hits + `` has no data!``. + """ + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + + detached = ht.detach() + _ = detached.prepare_for_saving() + + # Original must still be usable: dequantize, .device, repeated clone + _ = ht.dequantize() + _ = ht.device + + def test_detach_shares_underlying_buffers(self, input_tensor): + """Buffer tensors are shared (no GPU allocation) — only wrappers differ.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + detached = ht.detach() + + orig_row_buffers = ht._rowwise_storage.get_data_tensors() + new_row_buffers = detached._rowwise_storage.get_data_tensors() + if not isinstance(orig_row_buffers, tuple): + orig_row_buffers = (orig_row_buffers,) + new_row_buffers = (new_row_buffers,) + for a, b in zip(orig_row_buffers, new_row_buffers): + if a is None and b is None: + continue + assert a is b, "detach() must share buffer tensors, not copy them" + + @pytest.mark.parametrize("direction", ("rowwise", "columnwise")) + def test_detach_rejects_storage_only_sub_storage(self, input_tensor, direction): + """Storage-only children have no supported wrapper-copy protocol.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + getattr(hq, f"{direction}_quantizer").internal = True + ht = hq.quantize(input_tensor) + + with pytest.raises( + NotImplementedError, + match=rf"storage-only {direction} sub-storage", + ): + ht.detach() + + +@requires_fp8_and_nvfp4 +class TestHybridSaveRestore: + """Test prepare_for_saving / restore_from_saved round-trip.""" + + @pytest.fixture + def hybrid_tensor(self): + torch.manual_seed(42) + inp = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + return hq.quantize(inp) + + def test_save_restore_roundtrip(self, hybrid_tensor): + dq_before = hybrid_tensor.dequantize() + expected_metadata = { + "rowwise": _snapshot_storage_tensor_metadata( + hybrid_tensor.rowwise_sub_storage, clone=True + ), + "columnwise": _snapshot_storage_tensor_metadata( + hybrid_tensor.columnwise_sub_storage, clone=True + ), + } + buffers_before = _as_data_tensor_tuple(hybrid_tensor) + tensors, obj = hybrid_tensor.prepare_for_saving() + assert isinstance(tensors, list) + assert all(t is None or isinstance(t, torch.Tensor) for t in tensors) + + remainder = obj.restore_from_saved(tensors) + assert isinstance(remainder, list) + assert len(remainder) == 0 + + dq_after = hybrid_tensor.dequantize() + buffers_after = _as_data_tensor_tuple(hybrid_tensor) + assert len(buffers_after) == len(buffers_before) + for before, after in zip(buffers_before, buffers_after): + assert after is before, "Direct save/restore must reattach the exact saved buffer" + torch.testing.assert_close(dq_before, dq_after, rtol=0.0, atol=0.0) + for direction in ("rowwise", "columnwise"): + actual_metadata = _snapshot_storage_tensor_metadata( + getattr(hybrid_tensor, f"{direction}_sub_storage"), clone=False + ) + _assert_nested_state_exact( + actual_metadata, + expected_metadata[direction], + path=f"save/restore {direction}", + ) + + def test_save_clears_data(self, hybrid_tensor): + tensors, obj = hybrid_tensor.prepare_for_saving() + row_storage = hybrid_tensor.rowwise_sub_storage + row_data_tensors = row_storage.get_data_tensors() + if isinstance(row_data_tensors, tuple): + assert all(t is None for t in row_data_tensors) + else: + assert row_data_tensors is None + # Restore to clean up + obj.restore_from_saved(tensors) + + +@requires_fp8_and_nvfp4 +class TestHybridMakeEmpty: + """Test HybridQuantizer.make_empty().""" + + def test_make_empty_shape(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + shape = (128, 256) + empty = hq.make_empty(shape, dtype=torch.bfloat16, device="cuda") + assert isinstance(empty, HybridQuantizedTensor) + assert empty.shape == torch.Size(shape) + + def test_make_empty_dtype(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + shape = (128, 256) + empty = hq.make_empty(shape, dtype=torch.bfloat16, device="cuda") + assert empty.dtype == torch.bfloat16 + + def test_make_empty_has_sub_storages(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + shape = (128, 256) + empty = hq.make_empty(shape, dtype=torch.bfloat16, device="cuda") + assert empty.rowwise_sub_storage is not None + assert empty.columnwise_sub_storage is not None + + def test_make_empty_internal_returns_storage(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.internal = True + + empty = hq.make_empty((128, 256), dtype=torch.bfloat16, device="cuda") + + assert isinstance(empty, HybridQuantizedTensorStorage) + assert not isinstance(empty, HybridQuantizedTensor) + assert empty.size() == torch.Size((128, 256)) + assert empty.dequantize().dtype == torch.bfloat16 + + +@requires_fp8_and_nvfp4 +class TestHybridUsageFlagsRespected: + """``HybridQuantizer`` must skip directions whose parent usage flag is + False. Native quantizers honor ``rowwise_usage`` / ``columnwise_usage`` + inside the C++ kernel; hybrid sub-quantizers are pinned to one direction + in ``__init__``, so the parent's flags never reach C++ — the equivalent + skip lives in the Python composition layer. Modules call ``set_usage`` + extensively before each ``quantize`` (inference, output / grad_input + quantizers, AG paths), so honoring the flags avoids 2x quantization waste. + """ + + @pytest.fixture + def input_tensor(self): + torch.manual_seed(42) + return torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + + # ── quantize_impl ──────────────────────────────────────────── + + def test_quantize_rowwise_only(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.set_usage(rowwise=True, columnwise=False) + out = hq.quantize(input_tensor) + assert out.rowwise_sub_storage is not None + assert out.columnwise_sub_storage is None + + def test_quantize_columnwise_only(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.set_usage(rowwise=False, columnwise=True) + out = hq.quantize(input_tensor) + assert out.rowwise_sub_storage is None + assert out.columnwise_sub_storage is not None + + def test_quantize_both_false(self, input_tensor): + """``set_usage(False, False)`` mirrors ``update_usage(False, False)`` — + both produce an empty hybrid. No defensive assert (matches native).""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.set_usage(rowwise=False, columnwise=False) + out = hq.quantize(input_tensor) + assert out.rowwise_sub_storage is None + assert out.columnwise_sub_storage is None + + def test_quantize_both_true_default(self, input_tensor): + """Default state (both flags True) keeps both directions populated.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + out = hq.quantize(input_tensor) + assert out.rowwise_sub_storage is not None + assert out.columnwise_sub_storage is not None + + def test_quantize_internal_storage_rowwise_only(self, input_tensor): + """Internal storage path (used by FSDP2 / make_like flows) also + honors the gate.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.set_usage(rowwise=True, columnwise=False) + hq.internal = True + try: + out = hq.quantize(input_tensor) + assert isinstance(out, HybridQuantizedTensorStorage) + assert out.rowwise_sub_storage is not None + assert out.columnwise_sub_storage is None + finally: + hq.internal = False + + def test_quantize_flag_change_between_calls(self, input_tensor): + """A single quantizer can be re-used with different flags across + calls (which is exactly how modules use one ``input_quantizer`` / + ``weight_quantizer`` across forward / backward phases).""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + + hq.set_usage(rowwise=True, columnwise=False) + out_row = hq.quantize(input_tensor) + assert out_row.rowwise_sub_storage is not None + assert out_row.columnwise_sub_storage is None + + hq.set_usage(rowwise=False, columnwise=True) + out_col = hq.quantize(input_tensor) + assert out_col.rowwise_sub_storage is None + assert out_col.columnwise_sub_storage is not None + + hq.set_usage(rowwise=True, columnwise=True) + out_both = hq.quantize(input_tensor) + assert out_both.rowwise_sub_storage is not None + assert out_both.columnwise_sub_storage is not None + + # ── make_empty ─────────────────────────────────────────────── + + def test_make_empty_rowwise_only(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.set_usage(rowwise=True, columnwise=False) + empty = hq.make_empty((128, 256), dtype=torch.bfloat16, device="cuda") + assert empty.rowwise_sub_storage is not None + assert empty.columnwise_sub_storage is None + + def test_make_empty_columnwise_only(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.set_usage(rowwise=False, columnwise=True) + empty = hq.make_empty((128, 256), dtype=torch.bfloat16, device="cuda") + assert empty.rowwise_sub_storage is None + assert empty.columnwise_sub_storage is not None + + def test_make_empty_both_false(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.set_usage(rowwise=False, columnwise=False) + empty = hq.make_empty((128, 256), dtype=torch.bfloat16, device="cuda") + assert empty.rowwise_sub_storage is None + assert empty.columnwise_sub_storage is None + + # ── update_quantized ───────────────────────────────────────── + # + # Comparison strategy: snapshot raw data buffers via ``get_data_tensors()`` + # and compare bytes pre/post-update (same pattern as ``TestHybridClear``). + # Avoids per-format ``dequantize()`` limitations (NVFP4 columnwise raises + # NotImplementedError) and is a strictly stronger check — if the kernel + # writes, raw bytes differ regardless of whether dequant is reversible. + + @staticmethod + def _clone_data_tensors(sub_storage): + """Deep-clone the primary data buffers of a sub-storage.""" + if sub_storage is None: + return () + data = sub_storage.get_data_tensors() + if not isinstance(data, tuple): + data = (data,) + return tuple(t.clone() if t is not None else None for t in data) + + @staticmethod + def _assert_data_tensors_equal(snapshot, sub_storage): + """Assert sub-storage's current data buffers byte-match a prior snapshot.""" + assert sub_storage is not None + current = sub_storage.get_data_tensors() + if not isinstance(current, tuple): + current = (current,) + assert len(snapshot) == len( + current + ), f"Buffer count changed: {len(snapshot)} → {len(current)}" + for before, after in zip(snapshot, current): + if before is None: + assert after is None + continue + assert after is not None + torch.testing.assert_close(before, after, rtol=0, atol=0) + + @staticmethod + def _assert_data_tensors_differ(snapshot, sub_storage): + """Assert at least one buffer changed bytes vs the prior snapshot.""" + assert sub_storage is not None + current = sub_storage.get_data_tensors() + if not isinstance(current, tuple): + current = (current,) + any_changed = False + for before, after in zip(snapshot, current): + if before is None or after is None: + continue + if not torch.equal(before, after): + any_changed = True + break + assert any_changed, "Expected at least one data buffer to change but none did" + + def test_update_quantized_rowwise_only_preserves_columnwise_data(self, input_tensor): + """``update_quantized`` must not refresh a direction whose parent flag + is False, even if the dst storage has that direction allocated. + + Mirrors how native ``tex.quantize(src, quantizer, dst, noop_flag)`` + skips a direction when ``quantizer.rowwise_usage=False`` even if the + dst storage has that direction allocated. + """ + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + # Fully populate both directions + dst = hq.quantize(input_tensor) + # Snapshot the columnwise raw buffers before the targeted rowwise-only update + col_before = self._clone_data_tensors(dst._columnwise_storage) + + # Switch to rowwise-only refresh and feed a substantially different src + hq.set_usage(rowwise=True, columnwise=False) + new_src = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") * 100 + hq.update_quantized(new_src, dst) + + # Both sub-storage objects survive in-place; columnwise bytes untouched + self._assert_data_tensors_equal(col_before, dst._columnwise_storage) + + def test_update_quantized_columnwise_only_preserves_rowwise_data(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + dst = hq.quantize(input_tensor) + row_before = self._clone_data_tensors(dst._rowwise_storage) + + hq.set_usage(rowwise=False, columnwise=True) + new_src = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") * 100 + hq.update_quantized(new_src, dst) + + self._assert_data_tensors_equal(row_before, dst._rowwise_storage) + + def test_update_quantized_both_false_is_noop(self, input_tensor): + """``set_usage(False, False)`` then ``update_quantized`` must leave + both sub-storages' bytes untouched.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + dst = hq.quantize(input_tensor) + row_before = self._clone_data_tensors(dst._rowwise_storage) + col_before = self._clone_data_tensors(dst._columnwise_storage) + + hq.set_usage(rowwise=False, columnwise=False) + new_src = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") * 100 + hq.update_quantized(new_src, dst) + + self._assert_data_tensors_equal(row_before, dst._rowwise_storage) + self._assert_data_tensors_equal(col_before, dst._columnwise_storage) + + def test_update_quantized_actually_refreshes_requested(self, input_tensor): + """Sanity check: when the parent flag is True, the corresponding + sub-storage IS refreshed (otherwise the previous tests would pass + vacuously by not refreshing anything).""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + dst = hq.quantize(input_tensor) + row_before = self._clone_data_tensors(dst._rowwise_storage) + + hq.set_usage(rowwise=True, columnwise=False) + new_src = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") * 100 + hq.update_quantized(new_src, dst) + + # Rowwise bytes must differ — confirms update_quantized actually ran + self._assert_data_tensors_differ(row_before, dst._rowwise_storage) + + # ── te.Linear integration: inference path takes rowwise-only ─ + + def test_te_linear_inference_workspace_rowwise_only(self): + """``te.Linear`` forward under ``torch.no_grad()`` with a hybrid + ``CustomRecipe`` must produce a rowwise-only weight workspace. + ``linear.py:266-274`` sets ``weight_quantizer.set_usage(columnwise=False)`` + in inference; without the parent-flag gate, hybrid would still allocate + both directions. + """ + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + ) + torch.manual_seed(2026) + model = Linear(128, 256, bias=False, params_dtype=torch.bfloat16).cuda() + x = torch.randn(64, 128, dtype=torch.bfloat16, device="cuda") + + # is_first_microbatch=True forces the cache_name="weight" path + # (see linear.py:1631) so the hybrid workspace persists in + # model._fp8_workspaces and we can inspect its sub-storages. + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + _ = model(x, is_first_microbatch=True) + + ws = model._fp8_workspaces.get("weight") + assert isinstance( + ws, HybridQuantizedTensorStorage + ), f"Expected hybrid weight workspace, got {type(ws).__name__}" + assert ws.rowwise_sub_storage is not None, "Rowwise sub-storage must be populated for fprop" + assert ws.columnwise_sub_storage is None, ( + "Inference forward must produce rowwise-only hybrid weight workspace; " + "columnwise quantization should have been skipped per " + "weight_quantizer.set_usage(rowwise=True, columnwise=False)." + ) + + +@requires_fp8_and_nvfp4 +class TestHybridTorchDispatch: + """Test torch dispatch operations.""" + + @pytest.fixture + def hybrid_tensor(self): + torch.manual_seed(42) + inp = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + return hq.quantize(inp) + + def test_detach(self, hybrid_tensor): + detached = hybrid_tensor.detach() + assert isinstance(detached, HybridQuantizedTensor) + assert not detached.requires_grad + + def test_repr(self, hybrid_tensor): + r = repr(hybrid_tensor) + assert "HybridQuantizedTensor" in r + + +@requires_fp8_and_nvfp4 +class TestHybridGetDataTensors: + """Test get_data_tensors returns data from both sub-storages.""" + + def test_get_data_tensors(self): + torch.manual_seed(42) + inp = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(inp) + data_tensors = result.get_data_tensors() + row_tensors = _as_data_tensor_tuple(result.rowwise_sub_storage) + col_tensors = _as_data_tensor_tuple(result.columnwise_sub_storage) + + assert isinstance(data_tensors, tuple) + assert row_tensors and col_tensors + assert len(data_tensors) == len(row_tensors) + len(col_tensors) + assert all( + actual is expected for actual, expected in zip(data_tensors, row_tensors + col_tensors) + ), "Hybrid get_data_tensors must concatenate both sub-storages in direction order" + + +@requires_fp8_and_nvfp4 +class TestHybridDeviceAndSize: + """Test device and size properties.""" + + def test_device(self): + torch.manual_seed(42) + inp = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(inp) + assert result.device.type == "cuda" + + def test_size_from_storage(self): + torch.manual_seed(42) + inp = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.internal = True + result = hq.quantize(inp) + size = result.size() + assert size == torch.Size([128, 256]) + hq.internal = False + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridGemmBitwiseIdentical: + """Hybrid quantizer with same FP8 format in both directions must produce + bitwise-identical results to the vanilla Float8CurrentScaling recipe.""" + + def test_linear_fwd_bwd_matches_vanilla_fp8(self): + torch.manual_seed(123) + + in_features = 64 + out_features = 64 + batch = 32 + + model_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid.load_state_dict(model_ref.state_dict()) + + base_inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + inp_ref = base_inp.clone().detach().requires_grad_(True) + inp_hybrid = base_inp.clone().detach().requires_grad_(True) + + ref_recipe = recipe.Float8CurrentScaling() + with autocast(enabled=True, recipe=ref_recipe): + out_ref = model_ref(inp_ref) + loss_ref = out_ref.float().sum() + loss_ref.backward() + + def hybrid_fp8_factory(role): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("input", "weight", "output") + ): + return HybridQuantizer( + rowwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + ), + columnwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + ), + ) + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("grad_output", "grad_input") + ): + return Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, + device="cuda", + ) + return Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + ) + + hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_fp8_factory) + with autocast(enabled=True, recipe=hybrid_recipe): + out_hybrid = model_hybrid(inp_hybrid) + loss_hybrid = out_hybrid.float().sum() + loss_hybrid.backward() + + # Forward outputs must be bitwise identical + assert torch.equal( + out_ref, out_hybrid + ), f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" + + # Input gradients must be bitwise identical + assert inp_ref.grad is not None and inp_hybrid.grad is not None + assert torch.equal( + inp_ref.grad, inp_hybrid.grad + ), f"Input grad mismatch: max diff = {(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" + + # Parameter gradients must be bitwise identical + ref_params = dict(model_ref.named_parameters()) + hybrid_params = dict(model_hybrid.named_parameters()) + for name, p_ref in ref_params.items(): + p_hyb = hybrid_params[name] + assert ( + p_ref.grad is not None and p_hyb.grad is not None + ), f"Missing gradient for param '{name}'" + assert torch.equal(p_ref.grad, p_hyb.grad), ( + f"Param '{name}' grad mismatch: max diff = " + f"{(p_ref.grad - p_hyb.grad).abs().max().item()}" + ) + + +@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") +class TestHybridGemmBitwiseIdenticalMXFP8: + """Hybrid quantizer with MXFP8 in both directions must produce + bitwise-identical results to the vanilla MXFP8BlockScaling recipe.""" + + def test_linear_fwd_bwd_matches_vanilla_mxfp8(self): + torch.manual_seed(200) + + in_features, out_features, batch = 128, 128, 32 + + model_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid.load_state_dict(model_ref.state_dict()) + + base_inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + inp_ref = base_inp.clone().detach().requires_grad_(True) + inp_hybrid = base_inp.clone().detach().requires_grad_(True) + + ref_recipe = recipe.MXFP8BlockScaling() + with autocast(enabled=True, recipe=ref_recipe): + out_ref = model_ref(inp_ref) + out_ref.float().sum().backward() + + def hybrid_mxfp8_factory(role): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("grad_output", "grad_input") + ): + return MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + return HybridQuantizer( + rowwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + columnwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + ) + + hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_mxfp8_factory) + with autocast(enabled=True, recipe=hybrid_recipe): + out_hybrid = model_hybrid(inp_hybrid) + out_hybrid.float().sum().backward() + + assert torch.equal( + out_ref, out_hybrid + ), f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" + assert torch.equal( + inp_ref.grad, inp_hybrid.grad + ), f"Input grad mismatch: max diff = {(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" + for name, p_ref in dict(model_ref.named_parameters()).items(): + p_hyb = dict(model_hybrid.named_parameters())[name] + assert ( + p_ref.grad is not None and p_hyb.grad is not None + ), f"Missing gradient for param '{name}'" + assert torch.equal(p_ref.grad, p_hyb.grad), ( + f"Param '{name}' grad mismatch: max diff = " + f"{(p_ref.grad - p_hyb.grad).abs().max().item()}" + ) + + def test_dequantized_bwd_qfactory_save_original_input_matches_base_recipe_bitwise(self): + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + mxfp8_fwd_high_precision_bwd_factory, + ) + + torch.manual_seed(204) + in_features, out_features, batch = 128, 128, 32 + + model_ref = Linear( + in_features, + out_features, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=False, + ).cuda() + model_qfactory = Linear( + in_features, + out_features, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=True, + ).cuda() + model_qfactory.load_state_dict(model_ref.state_dict()) + + base_inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + inp_ref = base_inp.clone().detach().requires_grad_(True) + inp_qfactory = base_inp.clone().detach().requires_grad_(True) + + ref_recipe = recipe.MXFP8BlockScaling() + ref_recipe.backward_override = "dequantized" + with autocast(enabled=True, recipe=ref_recipe): + out_ref = model_ref(inp_ref) + + qfactory_recipe = recipe.CustomRecipe(qfactory=mxfp8_fwd_high_precision_bwd_factory) + with autocast(enabled=True, recipe=qfactory_recipe): + out_qfactory = model_qfactory(inp_qfactory) + + assert torch.equal( + out_ref, out_qfactory + ), f"Forward mismatch: max diff = {(out_ref - out_qfactory).abs().max().item()}" + + out_ref.float().sum().backward() + out_qfactory.float().sum().backward() + + assert inp_ref.grad is not None and inp_qfactory.grad is not None + assert torch.equal(inp_ref.grad, inp_qfactory.grad), ( + "Input grad mismatch: max diff = " + f"{(inp_ref.grad - inp_qfactory.grad).abs().max().item()}" + ) + for name, p_ref in dict(model_ref.named_parameters()).items(): + p_qfactory = dict(model_qfactory.named_parameters())[name] + assert ( + p_ref.grad is not None and p_qfactory.grad is not None + ), f"Missing gradient for param {name!r}" + assert torch.equal(p_ref.grad, p_qfactory.grad), ( + f"Param {name!r} grad mismatch: max diff = " + f"{(p_ref.grad - p_qfactory.grad).abs().max().item()}" + ) + + +@requires_fp8 +class TestCustomDPALocalRecipeCache: + """Custom-DPA native recipe labels track the quantizer rebuild.""" + + def test_inference_runs_once_per_quantizer_state_and_clears_stale_labels(self, monkeypatch): + from transformer_engine.pytorch.attention.dot_product_attention import ( + dot_product_attention as dpa_module, + ) + from transformer_engine.pytorch.module.base import TransformerEngineBaseModule + from transformer_engine.pytorch.quantization import FP8GlobalStateManager + + custom_recipe = recipe.CustomRecipe(qfactory=lambda _role: IdentityQuantizer()) + monkeypatch.setattr( + FP8GlobalStateManager, + "get_fp8_recipe", + classmethod(lambda _cls: custom_recipe), + ) + + state = [object()] + quantizer = [object()] + + def fake_base_init(module, num_gemms=1): # pylint: disable=unused-argument + module.fp8_meta["scaling_fwd"] = state[0] + module.quantizers["scaling_fwd"] = [quantizer[0]] + + monkeypatch.setattr( + TransformerEngineBaseModule, + "init_fp8_metadata", + fake_base_init, + ) + + inferred_labels = [recipe.MXFP8BlockScaling()] + mutated_recipe_labels = [recipe.MXFP8BlockScaling(fp8_mha=True)] + inference_results = iter((inferred_labels, mutated_recipe_labels, None)) + inference_calls = 0 + + def fake_infer(*_args, **_kwargs): + nonlocal inference_calls + inference_calls += 1 + return next(inference_results) + + monkeypatch.setattr(dpa_module, "_infer_custom_dpa_local_recipes", fake_infer) + + dpa = te.DotProductAttention( + num_attention_heads=2, + kv_channels=16, + attention_dropout=0.0, + ) + dpa.init_fp8_metadata() + assert dpa.fp8_meta["local_recipes"] is inferred_labels + dpa.init_fp8_metadata() + assert inference_calls == 1 + assert dpa.fp8_meta["local_recipes"] is inferred_labels + + # Native labels also copy these mutable fields from CustomRecipe. They + # must refresh even when the quantizer generation itself is unchanged. + custom_recipe.fp8_mha = True + dpa.init_fp8_metadata() + assert inference_calls == 2 + assert dpa.fp8_meta["local_recipes"] is mutated_recipe_labels + dpa.init_fp8_metadata() + assert inference_calls == 2 + assert dpa.fp8_meta["local_recipes"] is mutated_recipe_labels + + # A rebuilt recipe state/quantizer list invalidates the cache. If the + # new family has no native label, the old label must not survive. + state[0] = object() + quantizer[0] = object() + dpa.init_fp8_metadata() + assert inference_calls == 3 + assert "local_recipes" not in dpa.fp8_meta + dpa.init_fp8_metadata() + assert inference_calls == 3 + assert "local_recipes" not in dpa.fp8_meta + + @pytest.mark.parametrize( + "factory_name,expected", + [ + ("current_scaling_factory", (True, False)), + ("delayed_scaling_factory", (False, False)), + ], + ) + def test_qkv_capabilities_reuse_canonical_quantizer(self, factory_name, expected): + """Capability queries must not call qfactory outside recipe-state setup.""" + from transformer_engine.pytorch.custom_recipes import quantizer_factories + + base_qfactory = getattr(quantizer_factories, factory_name) + calls = [] + + def counting_qfactory(role): + calls.append(role) + # Model a factory with observable RNG state. An extra classification + # probe would consume another value and change subsequent results. + _ = torch.rand((), device="cuda") + return base_qfactory(role) + + custom_recipe = recipe.CustomRecipe( + qfactory=counting_qfactory, + fp8_dpa=True, + fp8_mha=True, + ) + dpa = te.DotProductAttention( + num_attention_heads=2, + kv_channels=16, + attention_dropout=0.0, + name="counted_dpa", + ).cuda() + + with autocast(enabled=True, recipe=custom_recipe): + first = dpa.get_qkv_quantization_capabilities() + canonical_qkv = dpa._qkv_capabilities_quantizer + calls_after_first = len(calls) + second = dpa.get_qkv_quantization_capabilities() + + assert first == expected + assert second == first + assert dpa._qkv_capabilities_quantizer is canonical_qkv + assert len(calls) == calls_after_first + + # A recipe-state rebuild creates a new canonical slot and must + # invalidate the capability cache automatically. + def rebuilt_qfactory(role): + return counting_qfactory(role) + + rebuilt_recipe = recipe.CustomRecipe( + qfactory=rebuilt_qfactory, + fp8_dpa=True, + fp8_mha=True, + ) + with autocast(enabled=True, recipe=rebuilt_recipe): + rebuilt = dpa.get_qkv_quantization_capabilities() + rebuilt_qkv = dpa._qkv_capabilities_quantizer + + assert rebuilt == first + assert rebuilt_qkv is not canonical_qkv + assert len(calls) > calls_after_first + + @pytest.mark.parametrize("fp8_mha", [False, True]) + def test_mha_forward_does_not_probe_qfactory(self, monkeypatch, fp8_mha): + """Repeated MHA forwards must not create discarded DPA quantizers.""" + from transformer_engine.pytorch.attention import multi_head_attention as mha_module + from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + current_scaling_factory, + delayed_scaling_factory, + ) + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + nvfp4_linear_fp8_dpa_factory, + ) + from transformer_engine.pytorch.utils import get_device_compute_capability + + cc = get_device_compute_capability() + if cc < (9, 0) or cc >= (12, 0): + pytest.skip(f"FP8 attention not supported on sm{cc[0] * 10 + cc[1]}") + + monkeypatch.setattr(mha_module, "_dpa_fp8_recipe", "") + monkeypatch.setenv("NVTE_UnfusedDPA_Emulate_FP8", "1") + calls = [] + + def valid_fp8_dpa_qfactory(role): + # Hopper FP8 attention supports DelayedScaling. Use it for the + # complete MHA recipe so that all requests share one coherent + # delayed-scaling state, including the DPA boundary tensors. + if cc < (10, 0): + return delayed_scaling_factory(role) + is_dpa = role is not None and role.module_type == "dpa" + is_dpa_boundary = ( + role is not None + and not role.module_type + and ("dpa_output" in role.name or "dpa_grad_input" in role.name) + ) + if is_dpa or is_dpa_boundary: + return nvfp4_linear_fp8_dpa_factory(role) + return current_scaling_factory(role) + + def counting_qfactory(role): + calls.append(role) + _ = torch.rand((), device="cuda") + return valid_fp8_dpa_qfactory(role) + + custom_recipe = recipe.CustomRecipe( + qfactory=counting_qfactory, + fp8_dpa=True, + fp8_mha=fp8_mha, + ) + model = te.MultiheadAttention( + hidden_size=128, + num_attention_heads=2, + kv_channels=64, + attention_dropout=0.0, + attn_mask_type="no_mask", + params_dtype=torch.bfloat16, + bias=False, + qkv_format="sbhd", + name="counted_mha", + ).cuda() + inp = torch.randn(128, 2, 128, device="cuda", dtype=torch.bfloat16) + + with torch.no_grad(), autocast(enabled=True, recipe=custom_recipe): + model(inp) + calls_after_first = len(calls) + model(inp) + + assert len(calls) == calls_after_first + + +@requires_fp8_and_nvfp4 +class TestAttentionFactoryNativeRecipeParity: + """Linear + DPA qfactories should match native DPA recipe switches bitwise.""" + + batch = 2 + seq_len = 128 + hidden_size = 128 + num_heads = 4 + kv_channels = hidden_size // num_heads + + class _LinearDPALinear(torch.nn.Module): + def __init__(self, hidden_size, num_heads, kv_channels): + super().__init__() + self.hidden_size = hidden_size + self.num_heads = num_heads + self.kv_channels = kv_channels + self.qkv_proj = Linear( + hidden_size, + 3 * hidden_size, + params_dtype=torch.bfloat16, + bias=False, + name="qkv", + ).cuda() + self.dpa = te.DotProductAttention( + num_heads, + kv_channels, + attention_dropout=0.0, + qkv_format="bshd", + name="core_attention", + ).cuda() + self.out_proj = Linear( + hidden_size, + hidden_size, + params_dtype=torch.bfloat16, + bias=False, + name="proj", + ).cuda() + + def forward(self, inp): + batch, seq_len, _ = inp.shape + qkv = self.qkv_proj(inp).view( + batch, + seq_len, + 3, + self.num_heads, + self.kv_channels, + ) + q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + attn_out = self.dpa(q, k, v, qkv_format="bshd").reshape( + batch, + seq_len, + self.hidden_size, + ) + return self.out_proj(attn_out) + + @staticmethod + def _set_native_dpa_recipe(monkeypatch, recipe_name): + from transformer_engine.pytorch.attention import multi_head_attention as mha_module + from transformer_engine.pytorch.attention.dot_product_attention import ( + dot_product_attention as dpa_module, + ) + + monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") + monkeypatch.setenv("NVTE_DPA_FP8_RECIPE", recipe_name) + monkeypatch.setenv("NVTE_DPA_FP8_FORMAT", "HYBRID") + monkeypatch.setenv("NVTE_DPA_FP8DS_AMAX_ALGO", "most_recent") + monkeypatch.setenv("NVTE_DPA_FP8DS_AMAX_HISTLEN", "1") + monkeypatch.setenv("NVTE_DPA_FP8DS_REDUCE_AMAX", "1") + monkeypatch.setattr(dpa_module, "_dpa_fp8_recipe", recipe_name) + monkeypatch.setattr(dpa_module, "_dpa_fp8_format", recipe.Format.HYBRID) + monkeypatch.setattr(dpa_module, "_dpa_fp8ds_amax_algo", "most_recent") + monkeypatch.setattr(dpa_module, "_dpa_fp8ds_amax_histlen", 1) + monkeypatch.setattr(dpa_module, "_dpa_fp8ds_reduce_amax", True) + monkeypatch.setattr(mha_module, "_dpa_fp8_recipe", recipe_name) + monkeypatch.setattr(mha_module, "_dpa_fp8_recipe_dpa", False) + monkeypatch.setattr(mha_module, "_dpa_fp8_recipe_mha", False) + + @staticmethod + def _clear_native_dpa_recipe(monkeypatch): + from transformer_engine.pytorch.attention import multi_head_attention as mha_module + from transformer_engine.pytorch.attention.dot_product_attention import ( + dot_product_attention as dpa_module, + ) + + monkeypatch.delenv("NVTE_DPA_FP8_RECIPE", raising=False) + monkeypatch.delenv("NVTE_DPA_FP8_FORMAT", raising=False) + monkeypatch.delenv("NVTE_DPA_FP8DS_AMAX_ALGO", raising=False) + monkeypatch.delenv("NVTE_DPA_FP8DS_AMAX_HISTLEN", raising=False) + monkeypatch.delenv("NVTE_DPA_FP8DS_REDUCE_AMAX", raising=False) + monkeypatch.setattr(dpa_module, "_dpa_fp8_recipe", "") + monkeypatch.setattr(mha_module, "_dpa_fp8_recipe", "") + monkeypatch.setattr(mha_module, "_dpa_fp8_recipe_dpa", False) + monkeypatch.setattr(mha_module, "_dpa_fp8_recipe_mha", False) + + @staticmethod + def _assert_equal(actual, expected, label): + assert torch.equal( + actual, expected + ), f"{label} mismatch: max diff = {(actual.float() - expected.float()).abs().max().item()}" + + def _run_model(self, model, inp, grad, fp8_recipe, seed): + run_inp = inp.clone().detach().requires_grad_(True) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + with autocast(enabled=True, recipe=fp8_recipe): + out = model(run_inp) + out.backward(grad) + local_recipes = [type(r).__name__ for r in model.dpa.fp8_meta.get("local_recipes", [])] + return ( + out.detach().clone(), + run_inp.grad.detach().clone(), + { + name: param.grad.detach().clone() + for name, param in model.named_parameters() + if param.grad is not None + }, + local_recipes, + ) + + @pytest.mark.parametrize( + "case_name,native_dpa_recipe,qfactory", + [ + ( + "fp8_dpa", + "Float8CurrentScaling", + nvfp4_linear_fp8_dpa_factory, + ), + ( + "mxfp8_dpa", + "MXFP8BlockScaling", + _nvfp4_linear_mxfp8_dpa_factory, + ), + ], + ) + def test_linear_dpa_linear_matches_native_env_recipe_bitwise( + self, + monkeypatch, + case_name, + native_dpa_recipe, + qfactory, + ): + if case_name == "mxfp8_dpa" and not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") + + from transformer_engine.pytorch.utils import get_device_compute_capability + + cc = get_device_compute_capability() + if cc < (9, 0) or cc >= (12, 0): + pytest.skip(f"FP8 attention not supported on sm{cc[0] * 10 + cc[1]}") + + self._set_native_dpa_recipe(monkeypatch, native_dpa_recipe) + + torch.manual_seed(2201) + model_native = self._LinearDPALinear( + self.hidden_size, + self.num_heads, + self.kv_channels, + ) + model_qfactory = self._LinearDPALinear( + self.hidden_size, + self.num_heads, + self.kv_channels, + ) + model_qfactory.load_state_dict(model_native.state_dict()) + + torch.manual_seed(2202) + base_inp = torch.randn( + self.batch, + self.seq_len, + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + ) + grad = torch.randn_like(base_inp) + + native_recipe = recipe.NVFP4BlockScaling(fp8_dpa=True) + qfactory_recipe = recipe.CustomRecipe(qfactory=qfactory, fp8_dpa=True) + + native_out, native_dx, native_grads, native_local_recipes = self._run_model( + model_native, + base_inp, + grad, + native_recipe, + seed=2203, + ) + self._clear_native_dpa_recipe(monkeypatch) + qfactory_out, qfactory_dx, qfactory_grads, qfactory_local_recipes = self._run_model( + model_qfactory, + base_inp, + grad, + qfactory_recipe, + seed=2203, + ) + + expected_local_recipes = ( + ["Float8CurrentScaling", "DelayedScaling"] + if native_dpa_recipe == "Float8CurrentScaling" + else ["MXFP8BlockScaling"] + ) + assert native_local_recipes == expected_local_recipes + assert qfactory_local_recipes == expected_local_recipes + + self._assert_equal(qfactory_out, native_out, f"{case_name} output") + self._assert_equal(qfactory_dx, native_dx, f"{case_name} input grad") + assert qfactory_grads.keys() == native_grads.keys() + for name, native_grad in native_grads.items(): + self._assert_equal( + qfactory_grads[name], + native_grad, + f"{case_name} param grad {name}", + ) + + def _run_mha_model(self, model, inp, grad, fp8_recipe, seed): + run_inp = inp.clone().detach().requires_grad_(True) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + with autocast(enabled=True, recipe=fp8_recipe): + out = model(run_inp, attn_mask_type="no_mask") + if isinstance(out, tuple): + out = out[0] + out.backward(grad) + return ( + out.detach().clone(), + run_inp.grad.detach().clone(), + { + name: param.grad.detach().clone() + for name, param in model.named_parameters() + if param.grad is not None + }, + ) + + @pytest.mark.parametrize( + "case_name,native_dpa_recipe,qfactory,expected_flags", + [ + ( + "fp8_dpa", + "Float8CurrentScaling", + nvfp4_linear_fp8_dpa_factory, + (False, False, False), + ), + ( + "mxfp8_dpa", + "MXFP8BlockScaling", + _nvfp4_linear_mxfp8_dpa_factory, + (False, False, False), + ), + ], + ) + def test_multihead_attention_matches_native_env_recipe_bitwise( + self, + monkeypatch, + case_name, + native_dpa_recipe, + qfactory, + expected_flags, + ): + if case_name == "mxfp8_dpa" and not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") + + from transformer_engine.pytorch.attention import multi_head_attention as mha_module + from transformer_engine.pytorch.utils import get_device_compute_capability + + cc = get_device_compute_capability() + if cc < (9, 0) or cc >= (12, 0): + pytest.skip(f"FP8 attention not supported on sm{cc[0] * 10 + cc[1]}") + + recorded_flags = [] + orig_update_roles = mha_module.MultiheadAttention._update_output_quantizer_roles + + def _recording_update_roles(self, qkv_fp8_output, proj_fp8_grad, dpa_fp8_output): + recorded_flags.append((qkv_fp8_output, dpa_fp8_output, proj_fp8_grad)) + return orig_update_roles(self, qkv_fp8_output, proj_fp8_grad, dpa_fp8_output) + + monkeypatch.setattr( + mha_module.MultiheadAttention, + "_update_output_quantizer_roles", + _recording_update_roles, + ) + + self._set_native_dpa_recipe(monkeypatch, native_dpa_recipe) + + torch.manual_seed(2301) + model_native = te.MultiheadAttention( + self.hidden_size, + self.num_heads, + kv_channels=self.kv_channels, + attention_dropout=0.0, + attn_mask_type="no_mask", + params_dtype=torch.bfloat16, + bias=False, + qkv_format="sbhd", + name="mha", + ).cuda() + model_qfactory = te.MultiheadAttention( + self.hidden_size, + self.num_heads, + kv_channels=self.kv_channels, + attention_dropout=0.0, + attn_mask_type="no_mask", + params_dtype=torch.bfloat16, + bias=False, + qkv_format="sbhd", + name="mha", + ).cuda() + model_qfactory.load_state_dict(model_native.state_dict()) + + torch.manual_seed(2302) + base_inp = torch.randn( + self.seq_len, + self.batch, + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + ) + grad = torch.randn_like(base_inp) + + native_recipe = recipe.NVFP4BlockScaling(fp8_dpa=True) + qfactory_recipe = recipe.CustomRecipe(qfactory=qfactory, fp8_dpa=True) + + native_out, native_dx, native_grads = self._run_mha_model( + model_native, + base_inp, + grad, + native_recipe, + seed=2303, + ) + native_flags = recorded_flags[-1] + self._clear_native_dpa_recipe(monkeypatch) + qfactory_out, qfactory_dx, qfactory_grads = self._run_mha_model( + model_qfactory, + base_inp, + grad, + qfactory_recipe, + seed=2303, + ) + qfactory_flags = recorded_flags[-1] + + assert native_flags == expected_flags + assert qfactory_flags == expected_flags + self._assert_equal(qfactory_out, native_out, f"{case_name} MHA output") + self._assert_equal(qfactory_dx, native_dx, f"{case_name} MHA input grad") + assert qfactory_grads.keys() == native_grads.keys() + for name, native_grad in native_grads.items(): + self._assert_equal( + qfactory_grads[name], + native_grad, + f"{case_name} MHA param grad {name}", + ) + + def test_update_output_quantizer_roles_wires_independent_boundaries(self): + from transformer_engine.pytorch.quantization import QuantizerRole + + model = te.MultiheadAttention( + self.hidden_size, + self.num_heads, + kv_channels=self.kv_channels, + attention_dropout=0.0, + attn_mask_type="no_mask", + params_dtype=torch.bfloat16, + bias=False, + qkv_format="sbhd", + name="mha", + ).cuda() + qkv = model.layernorm_qkv if model.input_layernorm else model.qkv + + expected_qkv = QuantizerRole( + module_type="dpa", + tensor_type="qkv", + name=model.core_attention.name or "", + ) + expected_do = QuantizerRole( + module_type="dpa", + tensor_type="do", + name=model.core_attention.name or "", + ) + expected_o = QuantizerRole( + module_type="linear", + tensor_type="input", + name=model.proj.name or "", + ) + expected_dqkv = QuantizerRole( + module_type="linear", + tensor_type="grad_output", + name=qkv.name or "", + ) + + def boundary_roles(): + return ( + qkv.output_quantizer_role, + model.proj.grad_input_quantizer_role, + model.core_attention.output_quantizer_role, + model.core_attention.grad_input_quantizer_role, + ) + + model._update_output_quantizer_roles(True, False, False) + assert boundary_roles() == (expected_qkv, None, None, None) + + model._update_output_quantizer_roles(False, True, False) + assert boundary_roles() == (None, expected_do, None, None) + + model._update_output_quantizer_roles(False, False, True) + assert boundary_roles() == (None, None, expected_o, expected_dqkv) + + model._update_output_quantizer_roles(False, False, False) + assert boundary_roles() == (None, None, None, None) + + def test_mxfp8_qfactory_uses_plain_bf16_mha_boundaries(self, monkeypatch): + """MXFP8 DPA stays internal; MHA boundary tensors remain plain BF16.""" + if not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") + + from transformer_engine.pytorch.attention.dot_product_attention import ( + backends as dpa_backends, + ) + from transformer_engine.pytorch.utils import get_device_compute_capability + + cc = get_device_compute_capability() + if cc < (9, 0) or cc >= (12, 0): + pytest.skip(f"FP8 attention not supported on sm{cc[0] * 10 + cc[1]}") + + self._clear_native_dpa_recipe(monkeypatch) + torch.manual_seed(2401) + model = te.MultiheadAttention( + self.hidden_size, + self.num_heads, + kv_channels=self.kv_channels, + attention_dropout=0.0, + attn_mask_type="no_mask", + params_dtype=torch.bfloat16, + bias=False, + qkv_format="sbhd", + name="mha", + ).cuda() + + boundary_tensors = {} + saved_dpa_tensors = {} + orig_prepare_for_saving = dpa_backends.prepare_for_saving + + def _record_prepare_for_saving(*tensors, **kwargs): + saved_dpa_tensors["fp8_o"] = tensors[3] + saved_dpa_tensors["f16_o"] = tensors[7] + return orig_prepare_for_saving(*tensors, **kwargs) + + monkeypatch.setattr(dpa_backends, "prepare_for_saving", _record_prepare_for_saving) + + def _record_grad(name): + def _hook(grad): + boundary_tensors[name] = grad + return grad + + return _hook + + def _record_dpa_inputs(_module, inputs, kwargs): + q, k, v = inputs[:3] + packed_qkv = kwargs.get("qkv_layer") + packed_kv = kwargs.get("kv_layer") + interleave_dim = kwargs.get("qkv_interleave_dim", -3) + + if packed_qkv is not None: + q, k, v = (packed_qkv.select(interleave_dim, i) for i in range(3)) + + def _record_packed_qkv_grad(grad): + for name, tensor in zip( + ("dq", "dk", "dv"), + (grad.select(interleave_dim, i) for i in range(3)), + ): + boundary_tensors[name] = tensor + + packed_qkv.register_hook(_record_packed_qkv_grad) + else: + q.register_hook(_record_grad("dq")) + if packed_kv is not None: + k, v = (packed_kv.select(interleave_dim, i) for i in range(2)) + + def _record_packed_kv_grad(grad): + boundary_tensors["dk"] = grad.select(interleave_dim, 0) + boundary_tensors["dv"] = grad.select(interleave_dim, 1) + + packed_kv.register_hook(_record_packed_kv_grad) + else: + k.register_hook(_record_grad("dk")) + v.register_hook(_record_grad("dv")) + + boundary_tensors.update(q=q, k=k, v=v) + + def _record_dpa_output(_module, _inputs, output): + boundary_tensors["o"] = output + output.register_hook(_record_grad("do")) + + def _record_projection_input(_module, inputs): + boundary_tensors["proj_input"] = inputs[0] + boundary_tensors["o_is_proj_input"] = boundary_tensors["o"] is inputs[0] + + handles = ( + model.core_attention.register_forward_pre_hook(_record_dpa_inputs, with_kwargs=True), + model.core_attention.register_forward_hook(_record_dpa_output), + model.proj.register_forward_pre_hook(_record_projection_input), + ) + + torch.manual_seed(2402) + inp = torch.randn( + self.seq_len, + self.batch, + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + ) + grad = torch.randn_like(inp) + qfactory_recipe = recipe.CustomRecipe( + qfactory=_nvfp4_linear_mxfp8_dpa_factory, + fp8_dpa=True, + ) + try: + self._run_mha_model(model, inp, grad, qfactory_recipe, seed=2403) + finally: + for handle in handles: + handle.remove() + + assert boundary_tensors["o_is_proj_input"] + assert saved_dpa_tensors["fp8_o"] is None + saved_o = saved_dpa_tensors["f16_o"] + assert type(saved_o) is torch.Tensor + assert saved_o.dtype is torch.bfloat16 + assert ( + saved_o.untyped_storage().data_ptr() + == boundary_tensors["o"].untyped_storage().data_ptr() + ) + + for name in ("q", "k", "v", "o", "proj_input", "dq", "dk", "dv", "do"): + tensor = boundary_tensors[name] + assert type(tensor) is torch.Tensor, f"{name} is {type(tensor)}" + assert tensor.dtype is torch.bfloat16, f"{name} has dtype {tensor.dtype}" + + +@pytest.mark.skipif(not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling) +class TestHybridGemmBitwiseIdenticalBlockFP8: + """Hybrid quantizer with Block FP8 in both directions must produce + bitwise-identical results to the vanilla Float8BlockScaling recipe.""" + + def test_linear_fwd_bwd_matches_vanilla_block_fp8(self): + torch.manual_seed(201) + + in_features, out_features, batch = 128, 128, 32 + + model_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid.load_state_dict(model_ref.state_dict()) + + base_inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + inp_ref = base_inp.clone().detach().requires_grad_(True) + inp_hybrid = base_inp.clone().detach().requires_grad_(True) + + ref_recipe = recipe.Float8BlockScaling() + with autocast(enabled=True, recipe=ref_recipe): + out_ref = model_ref(inp_ref) + out_ref.float().sum().backward() + + def hybrid_block_fp8_factory(role): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + is_weight = is_linear and role.tensor_type == "weight" + dim = 2 if is_weight else 1 + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=dim, + ) + return HybridQuantizer( + rowwise_quantizer=Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=dim, + ), + columnwise_quantizer=Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=dim, + ), + ) + + hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_block_fp8_factory) + with autocast(enabled=True, recipe=hybrid_recipe): + out_hybrid = model_hybrid(inp_hybrid) + out_hybrid.float().sum().backward() + + assert torch.equal( + out_ref, out_hybrid + ), f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" + assert torch.equal( + inp_ref.grad, inp_hybrid.grad + ), f"Input grad mismatch: max diff = {(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" + for name, p_ref in dict(model_ref.named_parameters()).items(): + p_hyb = dict(model_hybrid.named_parameters())[name] + assert ( + p_ref.grad is not None and p_hyb.grad is not None + ), f"Missing gradient for param '{name}'" + assert torch.equal(p_ref.grad, p_hyb.grad), ( + f"Param '{name}' grad mismatch: max diff = " + f"{(p_ref.grad - p_hyb.grad).abs().max().item()}" + ) + + +@pytest.mark.skipif( + not (fp8_available and nvfp4_available), + reason=f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", +) +class TestHybridGemmBitwiseIdenticalNVFP4: + """Same-format hybrid NVFP4 must match vanilla with seeded SR.""" + + def test_linear_fwd_bwd_matches_vanilla_nvfp4(self): + torch.manual_seed(202) + + in_features, out_features, batch = 128, 128, 32 + + model_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid.load_state_dict(model_ref.state_dict()) + + base_inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + inp_ref = base_inp.clone().detach().requires_grad_(True) + inp_hybrid = base_inp.clone().detach().requires_grad_(True) + + ref_recipe = recipe.NVFP4BlockScaling() + torch.manual_seed(1202) + torch.cuda.manual_seed_all(1202) + with autocast(enabled=True, recipe=ref_recipe): + out_ref = model_ref(inp_ref) + out_ref.float().sum().backward() + + def hybrid_nvfp4_factory(role): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type == "grad_output" + ): + return nvfp4_factory(role) + return HybridQuantizer( + rowwise_quantizer=nvfp4_factory(role), + columnwise_quantizer=nvfp4_factory(role), + ) + + hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_nvfp4_factory) + torch.manual_seed(1202) + torch.cuda.manual_seed_all(1202) + with autocast(enabled=True, recipe=hybrid_recipe): + out_hybrid = model_hybrid(inp_hybrid) + out_hybrid.float().sum().backward() + + assert torch.equal( + out_ref, out_hybrid + ), f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" + assert torch.equal( + inp_ref.grad, inp_hybrid.grad + ), f"Input grad mismatch: max diff = {(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" + for name, p_ref in dict(model_ref.named_parameters()).items(): + p_hyb = dict(model_hybrid.named_parameters())[name] + assert ( + p_ref.grad is not None and p_hyb.grad is not None + ), f"Missing gradient for param '{name}'" + assert torch.equal(p_ref.grad, p_hyb.grad), ( + f"Param '{name}' grad mismatch: max diff = " + f"{(p_ref.grad - p_hyb.grad).abs().max().item()}" + ) + + def test_linear_fwd_bwd_all_roles_hybrid(self): + """Exercise grad_output as a HybridQuantizer too.""" + torch.manual_seed(203) + + in_features, out_features, batch = 128, 128, 32 + + model_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid.load_state_dict(model_ref.state_dict()) + + base_inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + inp_ref = base_inp.clone().detach().requires_grad_(True) + inp_hybrid = base_inp.clone().detach().requires_grad_(True) + + ref_recipe = recipe.NVFP4BlockScaling() + torch.manual_seed(1203) + torch.cuda.manual_seed_all(1203) + with autocast(enabled=True, recipe=ref_recipe): + out_ref = model_ref(inp_ref) + out_ref.float().sum().backward() + + def hybrid_nvfp4_all_roles_factory(role): + return HybridQuantizer( + rowwise_quantizer=nvfp4_factory(role), + columnwise_quantizer=nvfp4_factory(role), + ) + + hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_nvfp4_all_roles_factory) + torch.manual_seed(1203) + torch.cuda.manual_seed_all(1203) + with autocast(enabled=True, recipe=hybrid_recipe): + out_hybrid = model_hybrid(inp_hybrid) + out_hybrid.float().sum().backward() + + assert torch.equal( + out_ref, out_hybrid + ), f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" + assert torch.equal( + inp_ref.grad, inp_hybrid.grad + ), f"Input grad mismatch: max diff = {(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" + for name, p_ref in dict(model_ref.named_parameters()).items(): + p_hyb = dict(model_hybrid.named_parameters())[name] + assert ( + p_ref.grad is not None and p_hyb.grad is not None + ), f"Missing gradient for param '{name}'" + assert torch.equal(p_ref.grad, p_hyb.grad), ( + f"Param '{name}' grad mismatch: max diff = " + f"{(p_ref.grad - p_hyb.grad).abs().max().item()}" + ) + + +@requires_fp8_and_nvfp4 +class TestHybridGemmMixedFormat: + """FP8 rowwise + NVFP4 columnwise through te.Linear forward+backward.""" + + def test_linear_fwd_bwd_fp8_row_nvfp4_col(self): + torch.manual_seed(42) + + in_features = 128 + out_features = 128 + batch = 32 + + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + inp = torch.randn( + batch, + in_features, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + + def mixed_factory(role): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("input", "weight") + ): + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("grad_output", "grad_input") + ): + return _make_nvfp4_quantizer() + return _make_fp8_quantizer() + + mixed_recipe = recipe.CustomRecipe(qfactory=mixed_factory) + + with autocast(enabled=True, recipe=mixed_recipe): + out = model(inp) + + assert out.shape == (batch, out_features) + assert out.dtype == torch.bfloat16 + assert not torch.isnan(out).any(), "Output contains NaN" + assert not torch.isinf(out).any(), "Output contains Inf" + + loss = out.float().sum() + loss.backward() + + assert inp.grad is not None, "Input gradient is None" + assert inp.grad.shape == inp.shape + assert not torch.isnan(inp.grad).any(), "Input gradient contains NaN" + assert not torch.isinf(inp.grad).any(), "Input gradient contains Inf" + + for name, p in model.named_parameters(): + assert p.grad is not None, f"Gradient for '{name}' is None" + assert not torch.isnan(p.grad).any(), f"Gradient for '{name}' contains NaN" + assert not torch.isinf(p.grad).any(), f"Gradient for '{name}' contains Inf" + + def test_numerical_sanity_against_bf16(self): + """Mixed-format output should be within reasonable tolerance of BF16 baseline.""" + torch.manual_seed(42) + + in_features = 128 + out_features = 128 + batch = 32 + + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + + # BF16 baseline (no quantization) + with torch.no_grad(): + out_bf16 = model(inp) + + def mixed_factory(role): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("input", "weight") + ): + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + return _make_fp8_quantizer() + + mixed_recipe = recipe.CustomRecipe(qfactory=mixed_factory) + with torch.no_grad(): + with autocast(enabled=True, recipe=mixed_recipe): + out_mixed = model(inp) + + # FP8/FP4 quantization introduces error, but the result should be + # in the same ballpark as BF16 + torch.testing.assert_close( + out_mixed.float(), + out_bf16.float(), + rtol=0.25, + atol=0.5, + ) + + +@requires_fp8_and_nvfp4 +class TestUnwrapTensor: + """Test GEMM input dispatch by rowwise or columnwise usage.""" + + @pytest.fixture + def hybrid_tensor(self): + torch.manual_seed(42) + inp = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + return hq.quantize(inp) + + def test_hybrid_rowwise_returns_rowwise_sub_storage(self, hybrid_tensor): + assert _unwrap_tensor(hybrid_tensor, "rowwise") is hybrid_tensor.rowwise_sub_storage + + def test_hybrid_columnwise_returns_columnwise_sub_storage(self, hybrid_tensor): + assert _unwrap_tensor(hybrid_tensor, "columnwise") is hybrid_tensor.columnwise_sub_storage + + @pytest.mark.parametrize( + "available_usage,missing_usage", + [("rowwise", "columnwise"), ("columnwise", "rowwise")], + ) + def test_missing_hybrid_representation_raises(self, available_usage, missing_usage): + inp = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + quantizer = _make_hybrid_quantizer_fp8_row_fp4_col() + quantizer.set_usage( + rowwise=available_usage == "rowwise", + columnwise=available_usage == "columnwise", + ) + hybrid_tensor = quantizer.quantize(inp) + + assert getattr(hybrid_tensor, f"{available_usage}_sub_storage") is not None + assert getattr(hybrid_tensor, f"{missing_usage}_sub_storage") is None + with pytest.raises( + RuntimeError, + match=rf"GEMM requested the {missing_usage} representation, but it is unavailable", + ): + _unwrap_tensor(hybrid_tensor, missing_usage) + + def test_rowwise_sub_storage_type(self, hybrid_tensor): + assert isinstance( + _unwrap_tensor(hybrid_tensor, "rowwise"), + (Float8TensorStorage, Float8Tensor), + ) + + def test_columnwise_sub_storage_type(self, hybrid_tensor): + assert isinstance( + _unwrap_tensor(hybrid_tensor, "columnwise"), + (NVFP4TensorStorage, NVFP4Tensor), + ) + + def test_non_hybrid_passthrough(self): + plain = torch.randn(4, 4, device="cuda") + for usage in ("rowwise", "columnwise"): + assert _unwrap_tensor(plain, usage) is plain + + def test_fp8_tensor_passthrough(self): + quantizer = _make_fp8_quantizer() + inp = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") + fp8 = quantizer.quantize(inp) + for usage in ("rowwise", "columnwise"): + assert _unwrap_tensor(fp8, usage) is fp8 + + def test_identity_falls_back_to_high_precision(self): + inp = torch.randn(4, 4, device="cuda") + identity = IdentityQuantizer()(inp) + + result = _unwrap_tensor(identity, "rowwise") + + assert isinstance(result, torch.Tensor) + assert not isinstance(result, QuantizedTensor) + assert torch.equal(result, inp) + + def test_custom_tensor_passthrough(self): + from transformer_engine.pytorch.custom_recipes.reference_current_scaling import ( + CurrentScalingTensorRef, + ) + + custom = CurrentScalingTensorRef() + assert _unwrap_tensor(custom, "rowwise") is custom + + def test_invalid_usage_raises(self): + with pytest.raises(ValueError, match="Unsupported GEMM tensor usage"): + _unwrap_tensor(torch.empty(0), "diagonal") + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridBiasGradient: + """Verify bias gradients are computed correctly with HybridQuantizer. + + tex.bgrad_quantize doesn't recognize HybridQuantizer, so the unfused + bgrad path is used instead. + """ + + def _make_uniform_hybrid_factory(self): + def factory(role): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("grad_output", "grad_input") + ): + return Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, + device="cuda", + ) + return HybridQuantizer( + rowwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + ), + columnwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + ), + ) + + return factory + + def test_bias_grad_matches_vanilla_fp8(self): + torch.manual_seed(456) + in_features, out_features, batch = 64, 64, 16 + + model_ref = Linear(in_features, out_features, bias=True, params_dtype=torch.bfloat16).cuda() + model_hybrid = Linear( + in_features, out_features, bias=True, params_dtype=torch.bfloat16 + ).cuda() + model_hybrid.load_state_dict(model_ref.state_dict()) + + base_inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + + # Reference + inp_ref = base_inp.clone().detach().requires_grad_(True) + with autocast(enabled=True, recipe=recipe.Float8CurrentScaling()): + out_ref = model_ref(inp_ref) + out_ref.float().sum().backward() + + # Hybrid + inp_hyb = base_inp.clone().detach().requires_grad_(True) + with autocast( + enabled=True, recipe=recipe.CustomRecipe(qfactory=self._make_uniform_hybrid_factory()) + ): + out_hyb = model_hybrid(inp_hyb) + out_hyb.float().sum().backward() + + ref_bias_grad = dict(model_ref.named_parameters())["bias"].grad + hyb_bias_grad = dict(model_hybrid.named_parameters())["bias"].grad + assert ref_bias_grad is not None and hyb_bias_grad is not None + assert torch.equal( + ref_bias_grad, hyb_bias_grad + ), f"Bias grad mismatch: max diff = {(ref_bias_grad - hyb_bias_grad).abs().max().item()}" + + def test_no_bias_fwd_bwd(self): + """Linear with bias=False skips bgrad_quantize entirely.""" + torch.manual_seed(42) + in_features, out_features, batch = 64, 64, 16 + + model = Linear(in_features, out_features, bias=False, params_dtype=torch.bfloat16).cuda() + inp = torch.randn( + batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) + + with autocast( + enabled=True, recipe=recipe.CustomRecipe(qfactory=self._make_uniform_hybrid_factory()) + ): + out = model(inp) + out.float().sum().backward() + + assert inp.grad is not None + assert not torch.isnan(inp.grad).any() + for name, p in model.named_parameters(): + assert p.grad is not None, f"Gradient for '{name}' is None" + + +@requires_fp8_and_nvfp4 +class TestHybridScalingModeCompatibility: + """cuBLAS requires matching scaling modes within a single GEMM. + + For hybrid quantization, this means the columnwise format for + linear_input/linear_weight must match the columnwise format for + linear_grad_output — otherwise the wgrad GEMM (NT layout) fails. + """ + + def test_matching_columnwise_formats_succeed(self): + """Both operands use NVFP4 columnwise → wgrad GEMM succeeds.""" + torch.manual_seed(42) + # NVFP4 GEMM requires dimensions ≥ 128 for cuBLAS support. + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + def factory(role): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("input", "weight") + ): + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("grad_output", "grad_input") + ): + return _make_nvfp4_quantizer() + return _make_fp8_quantizer() + + with autocast(enabled=True, recipe=recipe.CustomRecipe(qfactory=factory)): + out = model(inp) + out.float().sum().backward() + assert inp.grad is not None + + def test_mismatched_columnwise_formats_raise(self): + """NVFP4 input × FP8 grad_output columnwise → cuBLAS rejects.""" + torch.manual_seed(42) + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + def factory(role): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("input", "weight") + ): + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("grad_output", "grad_input") + ): + return Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, + device="cuda", + ) + return _make_fp8_quantizer() + + with autocast(enabled=True, recipe=recipe.CustomRecipe(qfactory=factory)): + out = model(inp) + with pytest.raises(RuntimeError, match="scaling_mode"): + out.float().sum().backward() + + +@requires_fp8_and_nvfp4 +class TestHybridReversedDirection: + """Reversed hybrid: NVFP4 rowwise (fprop) + FP8 columnwise (backward). + + Exercises NVFP4×NVFP4 in the fprop (TN) GEMM and FP8×FP8 in the + dgrad (NN) and wgrad (NT) GEMMs — the opposite of the primary + FP8-row/NVFP4-col configuration. + """ + + def test_nvfp4_row_fp8_col_forward_only(self): + """Forward (TN) with NVFP4×NVFP4 rowwise succeeds.""" + torch.manual_seed(99) + in_features, out_features, batch = 128, 128, 32 + + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + + def factory(role): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("input", "weight") + ): + return HybridQuantizer( + rowwise_quantizer=_make_nvfp4_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + return _make_nvfp4_quantizer() + + mixed_recipe = recipe.CustomRecipe(qfactory=factory) + with torch.no_grad(): + with autocast(enabled=True, recipe=mixed_recipe): + out = model(inp) + + assert out.shape == (batch, out_features) + assert not torch.isnan(out).any(), "Output contains NaN" + assert not torch.isinf(out).any(), "Output contains Inf" + + def test_nvfp4_row_fp8_col_full_fwd_bwd(self): + """Full fwd+bwd with NVFP4 rowwise (fprop) + FP8 columnwise (backward).""" + torch.manual_seed(99) + in_features, out_features, batch = 128, 128, 32 + + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + inp = torch.randn( + batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) + + def factory(role): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("input", "weight") + ): + return HybridQuantizer( + rowwise_quantizer=_make_nvfp4_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("grad_output", "grad_input") + ): + return _make_fp8_quantizer() + return _make_nvfp4_quantizer() + + mixed_recipe = recipe.CustomRecipe(qfactory=factory) + with autocast(enabled=True, recipe=mixed_recipe): + out = model(inp) + + assert out.shape == (batch, out_features) + assert not torch.isnan(out).any(), "Output contains NaN" + + loss = out.float().sum() + loss.backward() + + assert inp.grad is not None, "Input gradient is None" + assert not torch.isnan(inp.grad).any(), "Input gradient contains NaN" + for name, p in model.named_parameters(): + assert p.grad is not None, f"Gradient for '{name}' is None" + assert not torch.isnan(p.grad).any(), f"Gradient for '{name}' contains NaN" + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridMixedWithNonHybrid: + """Only one operand is hybrid; the other uses a plain TE quantizer. + + Exercises _unwrap_hybrid passthrough for the non-hybrid operand. + All roles must use compatible scaling modes for each GEMM: + fprop (TN): all rowwise formats must match + dgrad (NN): weight rowwise must match grad_output rowwise + wgrad (NT): input columnwise must match grad_output columnwise + """ + + @staticmethod + def _assert_native_fp8_parity(hybrid_role, *, seed): + torch.manual_seed(seed) + in_features, out_features, batch = 128, 128, 32 + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_native = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_native.load_state_dict(model_hybrid.state_dict()) + base_input = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + grad_output = torch.randn(batch, out_features, device="cuda", dtype=torch.bfloat16) + + def mixed_factory(role): + is_linear = role is not None and role.module_type in ( + "linear", + "grouped_linear", + ) + if is_linear and role.tensor_type == hybrid_role: + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda") + return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") + + hybrid_result = _run_linear_forward_backward( + model_hybrid, + base_input, + grad_output, + recipe.CustomRecipe(qfactory=mixed_factory), + seed=seed + 100, + ) + native_result = _run_linear_forward_backward( + model_native, + base_input, + grad_output, + recipe.Float8CurrentScaling(), + seed=seed + 100, + ) + _assert_linear_results_exact( + hybrid_result, + native_result, + output=True, + input_grad=True, + param_grads=True, + ) + + def test_hybrid_input_plain_weight_fwd_bwd(self): + """Input is hybrid (FP8 row / FP8 col), weight + grad_output plain FP8. + + Wgrad columnwise: FP8 (input.col) × FP8 (grad_output.col) → compatible. + """ + self._assert_native_fp8_parity("input", seed=77) + + def test_plain_input_hybrid_weight_fwd_bwd(self): + """Input is plain FP8, weight is hybrid (FP8 row / FP8 col).""" + self._assert_native_fp8_parity("weight", seed=88) + + +# --------------------------------------------------------------------------- +# Parametrized cross-format tests (stateless quantizers) +# --------------------------------------------------------------------------- + + +def _make_mxfp8_quantizer(*, rowwise=True, columnwise=True): + return MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=rowwise, + columnwise=columnwise, + ) + + +def _make_mxfp8_quantizer_e5m2(*, rowwise=True, columnwise=True): + return MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E5M2, + rowwise=rowwise, + columnwise=columnwise, + ) + + +def _make_block_quantizer(*, rowwise=True, columnwise=True): + return Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=rowwise, + columnwise=columnwise, + ) + + +def _make_block_quantizer_e5m2(*, rowwise=True, columnwise=True): + return Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E5M2, + rowwise=rowwise, + columnwise=columnwise, + ) + + +# (fwd_e4m3_factory, bwd_e5m2_factory, skip_condition, skip_reason) +_QUANTIZER_CONFIGS = { + "fp8_current": ( + _make_fp8_quantizer, + lambda **kw: Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda", **kw), + not fp8_available, + f"FP8: {reason_for_no_fp8}", + ), + "mxfp8": ( + _make_mxfp8_quantizer, + _make_mxfp8_quantizer_e5m2, + not mxfp8_available, + f"MXFP8: {reason_for_no_mxfp8}", + ), + "block_fp8": ( + _make_block_quantizer, + _make_block_quantizer_e5m2, + not fp8_block_scaling_available, + reason_for_no_fp8_block_scaling, + ), + "nvfp4": ( + _make_nvfp4_quantizer, + None, # NVFP4 has no E5M2 variant + not (fp8_available and nvfp4_available), + f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", + ), +} + + +def _build_cross_format_params(): + """Build parametrize list for all stateless cross-format hybrid combos.""" + combos = [ + ("fp8_current", "mxfp8"), + ("fp8_current", "nvfp4"), + ("fp8_current", "block_fp8"), + ("mxfp8", "fp8_current"), + ("mxfp8", "mxfp8"), + ("mxfp8", "nvfp4"), + ("mxfp8", "block_fp8"), + ("block_fp8", "fp8_current"), + ("block_fp8", "mxfp8"), + ("block_fp8", "nvfp4"), + ("block_fp8", "block_fp8"), + ("nvfp4", "fp8_current"), + ("nvfp4", "mxfp8"), + ("nvfp4", "block_fp8"), + ] + params = [] + for row, col in combos: + row_cfg = _QUANTIZER_CONFIGS[row] + col_cfg = _QUANTIZER_CONFIGS[col] + hw_skip = row_cfg[2] or col_cfg[2] + hw_reason = "; ".join( + filter(None, [row_cfg[3] if row_cfg[2] else "", col_cfg[3] if col_cfg[2] else ""]) + ) + marks = [] + if hw_skip: + marks.append(pytest.mark.skipif(True, reason=hw_reason or "N/A")) + if col == "fp8_current": + marks.append(_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8) + params.append(pytest.param(row, col, id=f"{row}_row_x_{col}_col", marks=marks)) + return params + + +def _set_quantization_test_seed(seed): + """Reset every RNG that a quantizer or GEMM helper may consume.""" + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def _run_linear_forward(model, base_input, fp8_recipe, *, seed): + """Run a deterministic forward pass and return a detached result.""" + _set_quantization_test_seed(seed) + with torch.no_grad(): + with autocast(enabled=True, recipe=fp8_recipe): + output = model(base_input) + return output.detach().clone() + + +def _run_linear_forward_backward(model, base_input, grad_output, fp8_recipe, *, seed): + """Run Linear with an external gradient and capture every numerical result.""" + model.zero_grad(set_to_none=True) + run_input = base_input.clone().detach().requires_grad_(True) + _set_quantization_test_seed(seed) + with autocast(enabled=True, recipe=fp8_recipe): + output = model(run_input) + output.backward(grad_output) + return ( + output.detach().clone(), + run_input.grad.detach().clone(), + { + name: param.grad.detach().clone() + for name, param in model.named_parameters() + if param.grad is not None + }, + ) + + +def _plain_linear_qfactory(operand_factory, grad_factory): + """Build a non-hybrid factory with role-correct operand and grad dtypes.""" + + def factory(role): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return _make_role_aware_quantizer(grad_factory, role) + return _make_role_aware_quantizer(operand_factory, role) + + return factory + + +def _assert_linear_results_exact(actual, expected, *, output, input_grad, param_grads): + """Compare selected Linear results with zero tolerance.""" + if output: + torch.testing.assert_close(actual[0], expected[0], rtol=0.0, atol=0.0) + if input_grad: + torch.testing.assert_close(actual[1], expected[1], rtol=0.0, atol=0.0) + if param_grads: + assert actual[2].keys() == expected[2].keys() + for name in actual[2]: + torch.testing.assert_close( + actual[2][name], + expected[2][name], + rtol=0.0, + atol=0.0, + msg=f"Parameter gradient {name!r} differs", + ) + + +class TestHybridCrossFormatParametrized: + """Parametrized fwd+bwd over all stateless quantizer cross-format pairs.""" + + @pytest.mark.parametrize("row_name,col_name", _build_cross_format_params()) + def test_fwd_bwd(self, row_name, col_name): + torch.manual_seed(42) + in_features, out_features, batch = 128, 128, 32 + + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_fprop_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_bwd_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_fprop_ref.load_state_dict(model_hybrid.state_dict()) + model_bwd_ref.load_state_dict(model_hybrid.state_dict()) + + base_input = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + grad_output = torch.randn(batch, out_features, device="cuda", dtype=torch.bfloat16) + + row_cfg = _QUANTIZER_CONFIGS[row_name] + col_cfg = _QUANTIZER_CONFIGS[col_name] + make_row_operand = row_cfg[0] + make_row_grad = row_cfg[1] if row_cfg[1] is not None else row_cfg[0] + make_col_operand = col_cfg[0] + make_col_grad = col_cfg[1] if col_cfg[1] is not None else col_cfg[0] + + def hybrid_factory(role): + is_linear = role is not None and role.module_type in ( + "linear", + "grouped_linear", + ) + if is_linear and role.tensor_type in ("input", "weight"): + return HybridQuantizer( + rowwise_quantizer=_make_role_aware_quantizer(make_row_operand, role), + columnwise_quantizer=_make_role_aware_quantizer(make_col_operand, role), + ) + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return _make_role_aware_quantizer(make_col_grad, role) + return _make_role_aware_quantizer(make_row_operand, role) + + hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_factory) + fprop_ref_recipe = recipe.CustomRecipe( + qfactory=_plain_linear_qfactory(make_row_operand, make_row_grad) + ) + bwd_ref_recipe = recipe.CustomRecipe( + qfactory=_plain_linear_qfactory(make_col_operand, make_col_grad) + ) + + hybrid_result = _run_linear_forward_backward( + model_hybrid, + base_input, + grad_output, + hybrid_recipe, + seed=1234, + ) + fprop_ref = _run_linear_forward( + model_fprop_ref, + base_input, + fprop_ref_recipe, + seed=1234, + ) + bwd_ref = _run_linear_forward_backward( + model_bwd_ref, + base_input, + grad_output, + bwd_ref_recipe, + seed=1234, + ) + + torch.testing.assert_close(hybrid_result[0], fprop_ref, rtol=0.0, atol=0.0) + _assert_linear_results_exact( + hybrid_result, + bwd_ref, + output=False, + input_grad=True, + param_grads=True, + ) + + +# --------------------------------------------------------------------------- +# CPU offload push/pop protocol (v2 OffloadableLayerState path) +# --------------------------------------------------------------------------- + + +class TestHybridCpuOffloadPushPop: + """Exercise the cpu_offload_v2 push/pop protocol on HybridQuantizedTensor. + + Uses :class:`OffloadableLayerState` directly — same pattern as + ``test_cpu_offloading.py::TestsOffloadableLayerState::test_general``. + Each test runs the full cycle: + + push → start_offload → release_activation_forward_gpu_memory + → start_reload → pop → release_all_memory + + The push path decomposes the hybrid via ``prepare_for_saving`` + (HybridQuantizedTensorStorage), recursively pushes each sub-storage + buffer, then reconstructs on pop via ``restore_from_saved``. Sub-buffers + below the 256K-element offload threshold (e.g. small block scales) are + returned unchanged; large data buffers round-trip through CPU. + """ + + # Hybrid tensor shape — each sub-storage primary buffer must exceed the + # cpu_offload _check_if_offload threshold (256K elements) so the path is + # actually exercised end-to-end. + _SHAPE = (1024, 1024) + + def _run_roundtrip(self, hybrid_tensor): + """Push → offload → release → reload → pop one hybrid tensor. + + Returns the reloaded tensor (a new HybridQuantizedTensor instance + reconstructed from the gathered-back buffers). + """ + from transformer_engine.pytorch.cpu_offload import OffloadableLayerState + + stream = torch.cuda.Stream() + state = OffloadableLayerState(offload_stream=stream) + + tid = state.push_tensor(hybrid_tensor) + state.start_offload() + state.release_activation_forward_gpu_memory() + state.start_reload() + reloaded = state.pop_tensor(tid) + torch.cuda.synchronize() + + try: + return reloaded + finally: + state.release_all_memory() + + @pytest.mark.parametrize("row_name,col_name", _build_cross_format_params()) + def test_push_pop_roundtrip(self, row_name, col_name): + """Dequantize-equivalence round-trip across the full 14-pair matrix.""" + torch.manual_seed(42) + inp = torch.randn(*self._SHAPE, dtype=torch.bfloat16, device="cuda") + + row_cfg = _QUANTIZER_CONFIGS[row_name] + col_cfg = _QUANTIZER_CONFIGS[col_name] + hq = HybridQuantizer( + rowwise_quantizer=row_cfg[0](), + columnwise_quantizer=col_cfg[0](), + ) + hybrid = hq.quantize(inp) + expected = hybrid.dequantize() + + reloaded = self._run_roundtrip(hybrid) + + _assert_hybrid_tensor_exact(reloaded, hybrid, context="CPU offload roundtrip") + assert isinstance(reloaded, HybridQuantizedTensor) + torch.testing.assert_close(reloaded.dequantize(), expected, rtol=0.0, atol=0.0) + + @requires_fp8_and_nvfp4 + def test_push_pop_preserves_sub_storage_types(self): + """Reconstructed hybrid preserves each sub-storage's concrete type.""" + torch.manual_seed(7) + inp = torch.randn(*self._SHAPE, dtype=torch.bfloat16, device="cuda") + + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hybrid = hq.quantize(inp) + row_type = type(hybrid.rowwise_sub_storage) + col_type = type(hybrid.columnwise_sub_storage) + + reloaded = self._run_roundtrip(hybrid) + + assert isinstance(reloaded.rowwise_sub_storage, row_type) + _assert_hybrid_tensor_exact(reloaded, hybrid, context="CPU offload storage types") + assert isinstance(reloaded.columnwise_sub_storage, col_type) + + @requires_fp8_and_nvfp4 + def test_push_pop_with_rowwise_only(self): + """Columnwise sub-storage dropped pre-push — roundtrip still works.""" + torch.manual_seed(11) + inp = torch.randn(*self._SHAPE, dtype=torch.bfloat16, device="cuda") + + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hybrid = hq.quantize(inp) + hybrid.update_usage(columnwise_usage=False) + assert hybrid.columnwise_sub_storage is None + expected = hybrid.dequantize() + + reloaded = self._run_roundtrip(hybrid) + + assert isinstance(reloaded, HybridQuantizedTensor) + assert reloaded.columnwise_sub_storage is None + _assert_hybrid_tensor_exact(reloaded, hybrid, context="CPU offload rowwise-only") + assert reloaded.rowwise_sub_storage is not None + torch.testing.assert_close(reloaded.dequantize(), expected, rtol=0.0, atol=0.0) + + @requires_fp8_and_nvfp4 + def test_push_pop_with_columnwise_only(self): + """Rowwise sub-storage dropped pre-push — roundtrip still works. + + Uses the reversed hybrid (NVFP4 rowwise + FP8 columnwise) so that + ``hybrid.dequantize()`` can fall through to the columnwise sub-storage. + ``HybridQuantizedTensorStorage.dequantize`` prefers rowwise and only + falls back to columnwise when rowwise is ``None``; NVFP4 does not yet + support columnwise-only dequantize, but Float8 does. + """ + torch.manual_seed(13) + inp = torch.randn(*self._SHAPE, dtype=torch.bfloat16, device="cuda") + + hq = _make_hybrid_quantizer_fp4_row_fp8_col() + hybrid = hq.quantize(inp) + hybrid.update_usage(rowwise_usage=False) + assert hybrid.rowwise_sub_storage is None + expected = hybrid.dequantize() + + reloaded = self._run_roundtrip(hybrid) + + assert isinstance(reloaded, HybridQuantizedTensor) + assert reloaded.rowwise_sub_storage is None + _assert_hybrid_tensor_exact(reloaded, hybrid, context="CPU offload columnwise-only") + assert reloaded.columnwise_sub_storage is not None + torch.testing.assert_close(reloaded.dequantize(), expected, rtol=0.0, atol=0.0) + + @requires_fp8_and_nvfp4 + def test_push_pop_roundtrip_does_not_leak_intermediate_buffers(self): + """After release_all_memory the offloader holds no hybrid buffers. + + Sanity check that the v2 cycle completes cleanly — no dangling CPU + pinned buffers left behind on a one-shot push/pop. + """ + from transformer_engine.pytorch.cpu_offload import OffloadableLayerState + + torch.manual_seed(17) + inp = torch.randn(*self._SHAPE, dtype=torch.bfloat16, device="cuda") + + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hybrid = hq.quantize(inp) + + stream = torch.cuda.Stream() + state = OffloadableLayerState(offload_stream=stream) + + tid = state.push_tensor(hybrid) + state.start_offload() + state.release_activation_forward_gpu_memory() + state.start_reload() + _ = state.pop_tensor(tid) + torch.cuda.synchronize() + state.release_all_memory() + + assert len(state.fwd_gpu_tensor_group.tensor_list) == 0 + assert len(state.cpu_tensor_group.tensor_list) == 0 + assert len(state.bwd_gpu_tensor_group.tensor_list) == 0 + assert state.state == "not_offloaded" + + +# --------------------------------------------------------------------------- +# 3-format hybrid: different quantization for fprop, dgrad, wgrad +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif( + not (fp8_available and mxfp8_available and nvfp4_available), + reason="Requires FP8 + MXFP8 + NVFP4", +) +class TestHybridThreeFormats: + """Three distinct formats: FormatA (fprop), FormatB (dgrad), FormatC (wgrad). + + Per-operand unwrap selects the correct sub-storage per GEMM: + fprop TN: weight.row(A) × input.row(A) → FormatA × FormatA + dgrad NN: weight.col(B) × grad_output.row(B) → FormatB × FormatB + wgrad NT: input.col(C) × grad_output.col(C) → FormatC × FormatC + + grad_output is itself hybrid (FormatB row + FormatC col) when B ≠ C. + """ + + @staticmethod + def _assert_three_format_routing( + make_fprop, make_dgrad, make_wgrad, *, seed, plain_grad_output=False + ): + in_features, out_features, batch = 128, 128, 32 + torch.manual_seed(seed) + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_fprop_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_dgrad_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_wgrad_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + state_dict = model_hybrid.state_dict() + model_fprop_ref.load_state_dict(state_dict) + model_dgrad_ref.load_state_dict(state_dict) + model_wgrad_ref.load_state_dict(state_dict) + + base_input = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + grad_output = torch.randn(batch, out_features, device="cuda", dtype=torch.bfloat16) + + def hybrid_factory(role): + is_linear = role is not None and role.module_type in ( + "linear", + "grouped_linear", + ) + if is_linear and role.tensor_type == "weight": + return HybridQuantizer( + rowwise_quantizer=make_fprop(), + columnwise_quantizer=make_dgrad(), + ) + if is_linear and role.tensor_type == "input": + return HybridQuantizer( + rowwise_quantizer=make_fprop(), + columnwise_quantizer=make_wgrad(), + ) + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + if plain_grad_output: + return make_dgrad() + return HybridQuantizer( + rowwise_quantizer=make_dgrad(), + columnwise_quantizer=make_wgrad(), + ) + return make_fprop() + + hybrid_result = _run_linear_forward_backward( + model_hybrid, + base_input, + grad_output, + recipe.CustomRecipe(qfactory=hybrid_factory), + seed=seed + 100, + ) + fprop_ref = _run_linear_forward( + model_fprop_ref, + base_input, + recipe.CustomRecipe(qfactory=_plain_linear_qfactory(make_fprop, make_fprop)), + seed=seed + 100, + ) + dgrad_ref = _run_linear_forward_backward( + model_dgrad_ref, + base_input, + grad_output, + recipe.CustomRecipe(qfactory=_plain_linear_qfactory(make_dgrad, make_dgrad)), + seed=seed + 100, + ) + wgrad_ref = _run_linear_forward_backward( + model_wgrad_ref, + base_input, + grad_output, + recipe.CustomRecipe(qfactory=_plain_linear_qfactory(make_wgrad, make_wgrad)), + seed=seed + 100, + ) + + torch.testing.assert_close(hybrid_result[0], fprop_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(hybrid_result[1], dgrad_ref[1], rtol=0.0, atol=0.0) + torch.testing.assert_close( + hybrid_result[2]["weight"], + wgrad_ref[2]["weight"], + rtol=0.0, + atol=0.0, + ) + torch.testing.assert_close( + hybrid_result[2]["bias"], + dgrad_ref[2]["bias"], + rtol=0.0, + atol=0.0, + ) + + def test_fp8_fprop_mxfp8_dgrad_nvfp4_wgrad(self): + """FP8 current (fprop) + MXFP8 (dgrad) + NVFP4 (wgrad).""" + self._assert_three_format_routing( + _make_fp8_quantizer, + _make_mxfp8_quantizer, + _make_nvfp4_quantizer, + seed=300, + ) + + def test_nvfp4_fprop_fp8_dgrad_mxfp8_wgrad(self): + """NVFP4 (fprop) + FP8 current (dgrad) + MXFP8 (wgrad).""" + self._assert_three_format_routing( + _make_nvfp4_quantizer, + _make_fp8_quantizer, + _make_mxfp8_quantizer, + seed=301, + ) + + def test_same_dgrad_wgrad_reduces_to_plain_grad(self): + """When dgrad format == wgrad format, grad_output can be a plain quantizer.""" + self._assert_three_format_routing( + _make_nvfp4_quantizer, + _make_mxfp8_quantizer, + _make_mxfp8_quantizer, + plain_grad_output=True, + seed=302, + ) + + +# --------------------------------------------------------------------------- +# All-modules test: hybrid quantization through every TE module type +# --------------------------------------------------------------------------- + + +def _make_hybrid_fp8_factory(): + """Factory returning HybridQuantizer(FP8 row + FP8 col) for fwd roles, + plain FP8 E5M2 for bwd roles.""" + + def factory(role): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight", "output"): + return HybridQuantizer( + rowwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + ), + columnwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + ), + ) + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, + device="cuda", + ) + return Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + ) + + return factory + + +@requires_fp8 +@pytest.mark.parametrize("norm_cls", (te.ops.LayerNorm, te.ops.RMSNorm)) +def test_fusible_norm_hybrid_output_falls_back_to_explicit_quantize(norm_cls): + """Unsupported fused Hybrid output is quantized by the following operation.""" + + def qfactory(_role): + return HybridQuantizer( + rowwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + ), + columnwise_quantizer=IdentityQuantizer(), + ) + + hidden_size = 128 + norm = norm_cls( + hidden_size, + device="cuda", + dtype=torch.bfloat16, + ) + forward = te.ops.Sequential(norm, te.ops.Quantize()) + inp = torch.randn( + 16, + hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + + with autocast(enabled=True, recipe=recipe.CustomRecipe(qfactory=qfactory)): + out = forward(inp) + + assert isinstance(out, HybridQuantizedTensor) + out.backward(torch.randn_like(inp)) + assert inp.grad is not None + assert norm.weight.grad is not None + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridAllModules: + """Hybrid quantization through all TE module types (not just Linear). + + Uses FP8 in both hybrid directions so the test validates module integration + without introducing cross-format scaling-mode concerns. + """ + + hidden_size = 128 + ffn_hidden_size = 128 + num_heads = 4 + batch = 16 + seq_len = 8 + + def _run_fwd_bwd(self, model, inp, *model_args, output_atol=0.0, param_atols=None): + """Compare same-format hybrid numerics against the native FP8 recipe.""" + import copy + + model_native = copy.deepcopy(model) + param_atols = {} if param_atols is None else param_atols + grad_output = torch.randn_like(inp) + + def run(run_model, run_recipe): + run_model.zero_grad(set_to_none=True) + run_input = inp.detach().clone().requires_grad_(True) + _set_quantization_test_seed(3370) + with autocast(enabled=True, recipe=run_recipe): + output = run_model(run_input, *model_args) + output.backward(grad_output) + return ( + output.detach().clone(), + run_input.grad.detach().clone(), + { + name: param.grad.detach().clone() + for name, param in run_model.named_parameters() + if param.grad is not None + }, + ) + + native_result = run(model_native, recipe.Float8CurrentScaling()) + hybrid_result = run( + model, + recipe.CustomRecipe(qfactory=_make_hybrid_fp8_factory()), + ) + + torch.testing.assert_close(hybrid_result[0], native_result[0], rtol=0.0, atol=output_atol) + torch.testing.assert_close(hybrid_result[1], native_result[1], rtol=0.0, atol=0.0) + assert hybrid_result[2].keys() == native_result[2].keys() + for name in hybrid_result[2]: + torch.testing.assert_close( + hybrid_result[2][name], + native_result[2][name], + rtol=0.0, + atol=param_atols.get(name, 0.0), + msg=f"Parameter gradient {name!r} differs", + ) + + def test_linear(self): + torch.manual_seed(500) + model = Linear( + self.hidden_size, + self.ffn_hidden_size, + params_dtype=torch.bfloat16, + ).cuda() + inp = torch.randn( + self.batch, + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + self._run_fwd_bwd(model, inp) + + def test_layernorm_linear(self): + torch.manual_seed(501) + model = LayerNormLinear( + self.hidden_size, + self.ffn_hidden_size, + params_dtype=torch.bfloat16, + ).cuda() + inp = torch.randn( + self.batch, + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + self._run_fwd_bwd(model, inp) + + def test_layernorm_mlp(self): + torch.manual_seed(502) + model = LayerNormMLP( + hidden_size=self.hidden_size, + ffn_hidden_size=self.ffn_hidden_size, + params_dtype=torch.bfloat16, + ).cuda() + inp = torch.randn( + self.batch, + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + # Native fuses LN+quantize while Hybrid takes the explicit quantize path. + # The only non-bitwise results are output (1 BF16 quantum here) and + # fc2_weight grad; all other gradients remain zero-tolerance checks. + self._run_fwd_bwd( + model, + inp, + output_atol=0.0009765625, + param_atols={"fc2_weight": 0.0234375}, + ) + + def test_grouped_linear(self): + torch.manual_seed(504) + num_gemms = 3 + model = GroupedLinear( + num_gemms, + self.hidden_size, + self.ffn_hidden_size, + params_dtype=torch.bfloat16, + ).cuda() + # Hopper FP8 grouped wgrad uses each expert's token count as a + # leading dimension, which cuBLAS requires to be divisible by 16. + m_splits = [16] * num_gemms + inp = torch.randn( + sum(m_splits), + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + + self._run_fwd_bwd(model, inp, m_splits) + + def test_transformer_layer(self): + torch.manual_seed(503) + model = TransformerLayer( + self.hidden_size, + self.ffn_hidden_size, + self.num_heads, + hidden_dropout=0.0, + attention_dropout=0.0, + fuse_qkv_params=True, + params_dtype=torch.bfloat16, + ).cuda() + inp = torch.randn( + self.seq_len, + self.batch, + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + # The LayerNormMLP submodule has the same fused-vs-explicit ordering. + # Keep dgrad and every other parameter exact; bound only the observed + # BF16 output and final projection-weight accumulation roundoff. + self._run_fwd_bwd( + model, + inp, + output_atol=0.015625, + param_atols={"layernorm_mlp.fc2_weight": 0.140625}, + ) + + +@requires_fp8 +class TestHybridGroupedLinearValidation: + """GroupedLinear generation-validation and split-dispatch coverage. + + Structural compatibility is validated once per real quantizer generation. + Steady-state dispatch reads the first expert after that uniformity check.""" + + @pytest.mark.parametrize( + "quantizers", + [ + pytest.param( + [_make_hybrid_quantizer_fp8_row_fp4_col() for _ in range(3)], + id="hybrid", + ), + pytest.param([_make_fp8_quantizer() for _ in range(3)], id="plain"), + pytest.param([None, None, None], id="none"), + ], + ) + def test_uniform_lists_validate(self, quantizers): + from transformer_engine.pytorch.module.grouped_linear import ( + _validate_grouped_quantizer_list, + ) + + _validate_grouped_quantizer_list(quantizers, operand_name="input") + + def test_plain_custom_quantizer_uses_python_split_fallback(self, monkeypatch): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + monkeypatch.setattr( + grouped_linear.tex, + "split_quantize", + lambda *args, **kwargs: pytest.fail("entered native split_quantize"), + ) + calls = [] + tensor = torch.randn((4, 8)) + quantizers = [_CountingPythonQuantizer(calls) for _ in range(2)] + + out = grouped_linear._split_quantize_non_hybrid( + tensor, + [2, 2], + quantizers, + tensor.dtype, + ) + + assert len(calls) == 2 + for actual, expected in zip(out, torch.split(tensor, [2, 2])): + torch.testing.assert_close(actual.dequantize(), expected, rtol=0.0, atol=0.0) + + @pytest.mark.parametrize("direction", ("rowwise", "columnwise")) + def test_hybrid_custom_child_uses_python_split_fallback( + self, + monkeypatch, + direction, + ): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + monkeypatch.setattr( + grouped_linear.tex, + "split_quantize", + lambda *args, **kwargs: pytest.fail("entered native split_quantize"), + ) + calls = [] + quantizers = [] + for _ in range(2): + rowwise = IdentityQuantizer() + columnwise = IdentityQuantizer() + if direction == "rowwise": + rowwise = _CountingPythonQuantizer(calls) + else: + columnwise = _CountingPythonQuantizer(calls) + quantizer = HybridQuantizer( + rowwise_quantizer=rowwise, + columnwise_quantizer=columnwise, + ) + quantizers.append(quantizer) + + out = grouped_linear._split_quantize_hybrid( + torch.randn((4, 8)), + [2, 2], + quantizers, + ) + + assert len(calls) == 2 + assert all(result.rowwise_sub_storage is not None for result in out) + assert all(result.columnwise_sub_storage is not None for result in out) + + def test_mixed_hybrid_and_plain_raises(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _validate_grouped_quantizer_list, + ) + + quantizers = [ + _make_hybrid_quantizer_fp8_row_fp4_col(), + _make_fp8_quantizer(), + _make_hybrid_quantizer_fp8_row_fp4_col(), + ] + with pytest.raises(ValueError, match="mix HybridQuantizer and non-hybrid"): + _validate_grouped_quantizer_list(quantizers, operand_name="input") + + def test_none_plus_hybrid_raises(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _validate_grouped_quantizer_list, + ) + + quantizers = [ + _make_hybrid_quantizer_fp8_row_fp4_col(), + None, + _make_hybrid_quantizer_fp8_row_fp4_col(), + ] + with pytest.raises(ValueError, match="mix None and concrete quantizers"): + _validate_grouped_quantizer_list(quantizers, operand_name="input") + + def test_mixed_identity_dtype_raises(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _validate_grouped_quantizer_list, + ) + + quantizers = [ + IdentityQuantizer(dtype=torch.bfloat16), + IdentityQuantizer(dtype=torch.float16), + ] + with pytest.raises(ValueError, match="incompatible plain backend configurations"): + _validate_grouped_quantizer_list(quantizers, operand_name="input") + + def test_distinct_delayed_scaling_state_is_allowed(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _validate_grouped_quantizer_list, + ) + + quantizers = [_make_delayed_quantizer(), _make_delayed_quantizer()] + quantizers[1].scale.fill_(2.0) + quantizers[1].amax.fill_(3.0) + _validate_grouped_quantizer_list(quantizers, operand_name="input") + + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + @pytest.mark.parametrize( + ("usage", "expected"), + [ + pytest.param((True, False), {"rowwise": True, "columnwise": False}, id="rowwise"), + pytest.param( + (False, True), + {"rowwise": False, "columnwise": True}, + id="columnwise", + marks=_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8, + ), + pytest.param( + (True, True), + {"rowwise": True, "columnwise": True}, + id="both", + marks=_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8, + ), + ], + ) + def test_hybrid_split_quantize_respects_parent_usage_flags(self, usage, expected): + from transformer_engine.pytorch.module.grouped_linear import ( + _split_quantize_hybrid, + ) + + tensor = torch.randn(32, 128, dtype=torch.bfloat16, device="cuda") + quantizers = [ + HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + for _ in range(2) + ] + for quantizer in quantizers: + quantizer.set_usage(rowwise=usage[0], columnwise=usage[1]) + + out = _split_quantize_hybrid(tensor, [16, 16], quantizers) + + assert [storage.get_usages() for storage in out] == [expected, expected] + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_columnwise_only_rowwise_dequantized_uses_transient_grouped_row(self, monkeypatch): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + real_split_quantize = grouped_linear.tex.split_quantize + calls = [] + + def tracked_split_quantize(tensor, m_splits, quantizers, **kwargs): + result = real_split_quantize(tensor, m_splits, quantizers, **kwargs) + calls.append((tensor, result)) + return result + + monkeypatch.setattr(grouped_linear.tex, "split_quantize", tracked_split_quantize) + tensor = torch.randn(32, 128, dtype=torch.bfloat16, device="cuda") + quantizers = [ + HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + columnwise_source="rowwise_dequantized", + ) + for _ in range(2) + ] + for quantizer in quantizers: + quantizer.set_usage(rowwise=False, columnwise=True) + + out = grouped_linear._split_quantize_hybrid(tensor, [16, 16], quantizers) + + assert len(calls) == 2 + expected_columnwise_source = torch.cat( + [result.dequantize(dtype=tensor.dtype) for result in calls[0][1]], dim=0 + ) + torch.testing.assert_close(calls[1][0], expected_columnwise_source, rtol=0.0, atol=0.0) + assert calls[1][0] is not tensor + assert all(storage.rowwise_sub_storage is None for storage in out) + assert all(storage.columnwise_sub_storage is not None for storage in out) + + @requires_nvfp4 + @pytest.mark.parametrize("m_splits", ([128, 128], [0, 128])) + def test_nvfp4_rowwise_dequantized_preserves_source_dtype(self, monkeypatch, m_splits): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + real_split_quantize = grouped_linear.tex.split_quantize + input_dtypes = [] + + def tracked_split_quantize(tensor, splits, quantizers, **kwargs): + input_dtypes.append(tensor.dtype) + return real_split_quantize(tensor, splits, quantizers, **kwargs) + + monkeypatch.setattr(grouped_linear.tex, "split_quantize", tracked_split_quantize) + tensor = torch.randn( + sum(m_splits), + 128, + dtype=torch.bfloat16, + device="cuda", + ) + quantizers = [ + HybridQuantizer( + rowwise_quantizer=_make_nvfp4_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + columnwise_source="rowwise_dequantized", + ) + for _ in m_splits + ] + + out = grouped_linear._split_quantize_hybrid(tensor, m_splits, quantizers) + + assert input_dtypes == [tensor.dtype, tensor.dtype] + assert len(out) == len(m_splits) + + def test_rowwise_only_skips_columnwise_quantization(self, monkeypatch): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + real_split_quantize = grouped_linear.tex.split_quantize + calls = [] + + def tracked_split_quantize(tensor, m_splits, quantizers, **kwargs): + calls.append(tensor) + return real_split_quantize(tensor, m_splits, quantizers, **kwargs) + + monkeypatch.setattr(grouped_linear.tex, "split_quantize", tracked_split_quantize) + tensor = torch.randn(32, 128, dtype=torch.bfloat16, device="cuda") + quantizers = [ + HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + columnwise_source="rowwise_dequantized", + ) + for _ in range(2) + ] + for quantizer in quantizers: + quantizer.set_usage(rowwise=True, columnwise=False) + + out = grouped_linear._split_quantize_hybrid(tensor, [16, 16], quantizers) + + assert calls == [tensor] + assert all(storage.rowwise_sub_storage is not None for storage in out) + assert all(storage.columnwise_sub_storage is None for storage in out) + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_original_source_preserves_two_bulk_call_fast_path(self, monkeypatch): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + real_split_quantize = grouped_linear.tex.split_quantize + calls = [] + + def tracked_split_quantize(tensor, m_splits, quantizers, **kwargs): + calls.append(tensor) + return real_split_quantize(tensor, m_splits, quantizers, **kwargs) + + monkeypatch.setattr(grouped_linear.tex, "split_quantize", tracked_split_quantize) + tensor = torch.randn(32, 128, dtype=torch.bfloat16, device="cuda") + quantizers = [ + HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + columnwise_source="original", + ) + for _ in range(2) + ] + + out = grouped_linear._split_quantize_hybrid(tensor, [16, 16], quantizers) + + assert calls == [tensor, tensor] + assert all(storage.rowwise_sub_storage is not None for storage in out) + assert all(storage.columnwise_sub_storage is not None for storage in out) + + def test_validation_rejects_mixed_columnwise_source_policies(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _validate_grouped_quantizer_list, + ) + + quantizers = [ + HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + columnwise_source=source, + ) + for source in ("original", "rowwise_dequantized") + ] + + with pytest.raises(ValueError, match="mixed columnwise source policies"): + _validate_grouped_quantizer_list(quantizers, operand_name="input") + + def test_validation_rejects_same_family_config_mismatch(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _validate_grouped_quantizer_list, + ) + + quantizers = [_make_fp8_quantizer(), _make_fp8_quantizer()] + quantizers[1].force_pow_2_scales = True + + with pytest.raises( + ValueError, + match="incompatible plain backend configurations", + ): + _validate_grouped_quantizer_list(quantizers, operand_name="input") + + def test_validation_runs_only_with_quantizer_generation(self, monkeypatch): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + def make_qfactory(columnwise_source): + def qfactory(_role): + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + columnwise_source=columnwise_source, + ) + + return qfactory + + model = GroupedLinear(2, 128, 128, bias=False, params_dtype=torch.bfloat16).cuda() + tensor = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + m_splits = torch.tensor([64, 64], dtype=torch.int64) + original_recipe = recipe.CustomRecipe(qfactory=make_qfactory("original")) + + real_validate = grouped_linear._validate_grouped_quantizer_list + validation_calls = [] + + def tracked_validate(quantizers, *, operand_name="operand"): + validation_calls.append((operand_name, id(quantizers[0]))) + return real_validate(quantizers, operand_name=operand_name) + + monkeypatch.setattr( + grouped_linear, + "_validate_grouped_quantizer_list", + tracked_validate, + ) + + with torch.no_grad(), autocast(enabled=True, recipe=original_recipe): + model(tensor, m_splits) + first_call_count = len(validation_calls) + first_generation = model._validated_quantizer_generations["scaling_fwd"] + assert first_call_count > 0 + + with torch.no_grad(), autocast(enabled=True, recipe=original_recipe): + model(tensor, m_splits) + assert len(validation_calls) == first_call_count + assert model._validated_quantizer_generations["scaling_fwd"] is first_generation + + rebuilt_recipe = recipe.CustomRecipe(qfactory=make_qfactory("rowwise_dequantized")) + with torch.no_grad(), autocast(enabled=True, recipe=rebuilt_recipe): + model(tensor, m_splits) + rebuilt_generation = model._validated_quantizer_generations["scaling_fwd"] + assert len(validation_calls) > first_call_count + assert rebuilt_generation is not first_generation + assert rebuilt_generation[0].columnwise_source == "rowwise_dequantized" + + input_count = 0 + + def mixed_source_qfactory(role): + nonlocal input_count + source = "original" + if role is not None and role.tensor_type == "input": + source = "original" if input_count == 0 else "rowwise_dequantized" + input_count += 1 + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + columnwise_source=source, + ) + + mixed_recipe = recipe.CustomRecipe(qfactory=mixed_source_qfactory) + # A failed generation is never marked validated. Base metadata can then + # early-return on retry, so the O(1) guard must validate it again. + for _ in range(2): + with pytest.raises(ValueError, match="mixed columnwise source policies"): + with torch.no_grad(), autocast(enabled=True, recipe=mixed_recipe): + model(tensor, m_splits) + assert model._validated_quantizer_generations["scaling_fwd"] is rebuilt_generation + + # Stale invalid recipe metadata must not affect the non-quantized path. + with torch.no_grad(): + model(tensor, m_splits) + + @requires_fp8_and_nvfp4 + def test_hybrid_split_quantize_honors_rowwise_dequantized_source(self): + """NVFP4 column data must derive from the actual grouped row result.""" + from transformer_engine.pytorch.module.grouped_linear import ( + _split_quantize_hybrid, + ) + + torch.manual_seed(3598) + # NVFP4 grouped split-quantize requires each M split to be a multiple + # of 64. + tensor = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + + def make_quantizer(): + return HybridQuantizer( + rowwise_quantizer=_make_nvfp4_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + columnwise_source="rowwise_dequantized", + ) + + quantizers = [make_quantizer(), make_quantizer()] + actual = _split_quantize_hybrid( + tensor, + [64, 64], + quantizers, + ) + + for index, (actual_part, quantizer) in enumerate(zip(actual, quantizers)): + expected_columnwise = quantizer.columnwise_quantizer.quantize( + actual_part.rowwise_sub_storage.dequantize() + ) + _assert_storage_data_exact( + actual_part.columnwise_sub_storage, + expected_columnwise, + context=f"GroupedLinear split {index} columnwise provenance", + ) + + +# =========================================================================== +# Quantized Parameters (quantized_model_init) tests for hybrid quantization +# =========================================================================== + +# --------------------------------------------------------------------------- +# 1. quantized_model_init: model creation and parameter type verification +# --------------------------------------------------------------------------- + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridQuantizedModelInit: + """Verify that quantized_model_init with a hybrid CustomRecipe produces + HybridQuantizedTensor parameters.""" + + def _hybrid_fp8_recipe(self): + return _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + + def test_linear_weight_is_hybrid_quantized_tensor(self): + """model.weight should be a HybridQuantizedTensor after quantized_model_init.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + weight = model.weight + assert isinstance( + weight, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(weight).__name__}" + assert isinstance( + weight, QuantizedTensor + ), "HybridQuantizedTensor should be a QuantizedTensor subclass" + + def test_linear_weight_has_both_sub_storages(self): + """Quantized param should have rowwise and columnwise sub-storages.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + weight = model.weight + assert weight.rowwise_sub_storage is not None, "Missing rowwise sub-storage" + assert weight.columnwise_sub_storage is not None, "Missing columnwise sub-storage" + + def test_linear_weight_shape_preserved(self): + """Quantized param should retain its logical shape.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 256, params_dtype=torch.bfloat16).cuda() + + assert model.weight.shape == torch.Size([256, 128]) + + def test_linear_bias_stays_bf16(self): + """Bias should remain BF16 (not quantized).""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, bias=True, params_dtype=torch.bfloat16).cuda() + + assert not isinstance(model.bias, QuantizedTensor), "Bias should not be a QuantizedTensor" + assert model.bias.dtype == torch.bfloat16 + + def test_layernorm_linear_weight_is_hybrid(self): + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = LayerNormLinear(128, 128, params_dtype=torch.bfloat16).cuda() + + assert isinstance(model.weight, HybridQuantizedTensor) + + def test_dequantize_close_to_original(self): + """Dequantized hybrid param should be close to the BF16 init values.""" + hybrid_recipe = self._hybrid_fp8_recipe() + + # Create a non-quantized reference + torch.manual_seed(42) + ref_model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + ref_weight = ref_model.weight.detach().clone() + + # Create quantized model with the same seed + torch.manual_seed(42) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + dq_weight = model.weight.dequantize() + torch.testing.assert_close(dq_weight.float(), ref_weight.float(), rtol=0.125, atol=0.1) + + def test_preserve_high_precision_init_val(self): + """preserve_high_precision_init_val should store original BF16 on CPU.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init( + enabled=True, + recipe=hybrid_recipe, + preserve_high_precision_init_val=True, + ): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + weight = model.weight + assert isinstance(weight, HybridQuantizedTensor) + assert hasattr(weight, "get_high_precision_init_val") + hp_val = weight.get_high_precision_init_val() + assert hp_val is not None, "High-precision init val should be stored" + assert hp_val.device.type == "cpu" + assert hp_val.shape == weight.shape + + +# --------------------------------------------------------------------------- +# 2. get_weight_workspace cache invalidation for hybrid +# --------------------------------------------------------------------------- + + +@requires_fp8 +class TestHybridWeightWorkspaceCache: + """Test that get_weight_workspace handles HybridQuantizedTensorStorage + correctly for the quantized-params early-return path and the BF16 cache path.""" + + def _hybrid_fp8_recipe(self): + return _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_quantized_param_skips_workspace(self): + """When weight is already a HybridQuantizedTensor (quantized params), + get_weight_workspace should return it directly without creating a workspace.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + with autocast(enabled=True, recipe=hybrid_recipe): + out = model(inp, is_first_microbatch=True) + + assert out.shape == (32, 128) + assert "weight" not in model._fp8_workspaces + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_bf16_weight_creates_hybrid_workspace(self): + """When weight is BF16 and recipe produces HybridQuantizer, the workspace + should be a HybridQuantizedTensor.""" + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + hybrid_recipe = self._hybrid_fp8_recipe() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + with autocast(enabled=True, recipe=hybrid_recipe): + out = model(inp, is_first_microbatch=True) + + assert out.shape == (32, 128) + workspace = model._fp8_workspaces.get("weight") + assert isinstance(workspace, HybridQuantizedTensorStorage) + assert workspace.rowwise_sub_storage is not None + assert workspace.columnwise_sub_storage is not None + + def test_workspace_cache_reuse_across_microbatches(self): + """Cached hybrid workspace should be reused on 2nd+ microbatches.""" + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + hybrid_recipe = self._hybrid_fp8_recipe() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + with autocast(enabled=True, recipe=hybrid_recipe): + with torch.no_grad(): + out1 = model(inp, is_first_microbatch=True) + workspace = model._fp8_workspaces["weight"] + buffers = _as_data_tensor_tuple(workspace) + out2 = model(inp, is_first_microbatch=False) + + assert isinstance(workspace, HybridQuantizedTensorStorage) + assert model._fp8_workspaces["weight"] is workspace + current_buffers = _as_data_tensor_tuple(workspace) + assert len(current_buffers) == len(buffers) + assert all(current is cached for current, cached in zip(current_buffers, buffers)) + torch.testing.assert_close(out1, out2, rtol=0.0, atol=0.0) + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_workspace_cache_invalidation_on_usage_change(self): + """If usage requirements change (e.g. inference→training), the cache + should be invalidated and a fresh workspace created.""" + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + hybrid_recipe = self._hybrid_fp8_recipe() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + # First pass: inference (no columnwise needed) + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + model(inp, is_first_microbatch=True) + inference_workspace = model._fp8_workspaces["weight"] + assert isinstance(inference_workspace, HybridQuantizedTensorStorage) + assert inference_workspace.rowwise_sub_storage is not None + assert inference_workspace.columnwise_sub_storage is None + + # Second pass: training (columnwise now needed for backward) + with autocast(enabled=True, recipe=hybrid_recipe): + out_train = model(inp, is_first_microbatch=True) + training_workspace = model._fp8_workspaces["weight"] + + assert isinstance(training_workspace, HybridQuantizedTensorStorage) + assert training_workspace is not inference_workspace + assert training_workspace.rowwise_sub_storage is not None + assert training_workspace.columnwise_sub_storage is not None + + loss = out_train.float().sum() + loss.backward() + + assert inp.grad is not None + + +# --------------------------------------------------------------------------- +# 3. _update_weight_quantizers for hybrid +# --------------------------------------------------------------------------- + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridUpdateWeightQuantizers: + """Test that quantizer refresh propagates correctly to hybrid sub-quantizers.""" + + def _hybrid_fp8_recipe(self): + return _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + + def test_quantized_param_survives_multiple_forward_passes(self): + """Weight should remain a HybridQuantizedTensor across multiple forward passes, + each of which triggers init_fp8_metadata → potential quantizer updates.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + for i in range(3): + inp_i = inp.detach().clone().requires_grad_(True) + with autocast(enabled=True, recipe=hybrid_recipe): + out = model(inp_i) + out.float().sum().backward() + assert not torch.isnan(out).any(), f"NaN at iteration {i}" + assert inp_i.grad is not None, f"No input grad at iteration {i}" + + assert isinstance( + model.weight, HybridQuantizedTensor + ), "Weight lost HybridQuantizedTensor type after multiple passes" + + +# --------------------------------------------------------------------------- +# 3b. quantize_master_weights + post_all_gather_processing for hybrid params +# +# Covers the supported (same-format) cases and the rejected (cross-format, +# missing sub-storage, unsupported sub-quantizer) cases. The supported subset +# is the first incremental hybrid integration with the distributed-optimizer +# quantized-param all-gather flow. Cross-format support is deferred to +# follow-up #3158; the tests below pin the NotImplementedError contract so the +# rejection messaging stays clear as the feature evolves. +# --------------------------------------------------------------------------- + + +def _ensure_single_rank_dp_group(): + """Return a single-rank NCCL process group for hybrid quantize_master_weights + tests. Mirrors the local-pytest setup in + `tests/pytorch/distributed/test_cast_master_weights_to_fp8.py` so we can call + `torch.distributed.all_reduce` against a trivial group from inside the + per-format helpers. The group is created lazily on first call and reused + across tests within the same pytest process. + """ + # pylint: disable=import-outside-toplevel + import tempfile + import pathlib + + if not torch.distributed.is_initialized(): + torch.cuda.set_device(0) + with tempfile.NamedTemporaryFile(delete=False) as f: + rendezvous_file = pathlib.Path(f.name) + torch.distributed.init_process_group( + backend="nccl", + init_method=rendezvous_file.resolve().as_uri(), + rank=0, + world_size=1, + ) + return torch.distributed.GroupMember.WORLD + + +def _hybrid_recipe_fp8_current(): + """Same-format Float8CurrentScaling on both directions (supported).""" + return _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + grad_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda"), + ) + + +def _make_delayed_quantizer(fp8_dtype=None, *, rowwise=True, columnwise=True): + """Construct a ``Float8Quantizer`` (delayed scaling) with locally-allocated + scale/amax buffers for single-shot unit tests. + + The full delayed-scaling lifecycle (``FP8GlobalStateManager`` updating + ``amax_history`` -> ``scale`` across iterations) is out of scope here; for + ``quantize_master_weights`` we only need the helper to read/write + ``quantizer.amax`` / ``quantizer.scale`` / ``model_weight._scale_inv``, + which works with any pair of 1-element float32 tensors. Initial scale=1.0 + and amax=0.0 mirror the cold-start state ``FP8GlobalStateManager`` would + initialize for the first iteration. + """ + if fp8_dtype is None: + fp8_dtype = tex.DType.kFloat8E4M3 + return Float8Quantizer( + scale=torch.ones(1, dtype=torch.float32, device="cuda"), + amax=torch.zeros(1, dtype=torch.float32, device="cuda"), + fp8_dtype=fp8_dtype, + rowwise=rowwise, + columnwise=columnwise, + ) + + +def _hybrid_recipe_fp8_delayed(): + """Same-format Float8 delayed scaling on both directions (supported).""" + return _hybrid_custom_recipe( + row_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E4M3), + col_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E4M3), + grad_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E5M2), + ) + + +def _hybrid_recipe_fp8_delayed_row_current_col(): + """Cross-format per-tensor Float8: delayed rowwise + current columnwise. + + Routed per-direction: row sub-storage -> delayed bucket, col sub-storage + -> current bucket. The two helpers run independently (no shared state), + so each direction's scale is computed via its own scaling lifecycle. + """ + return _hybrid_custom_recipe( + row_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E4M3), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + # grad_factory matches the columnwise direction so the wgrad GEMM's + # grad_output sub-quantizer pairs with the input/weight col format. + grad_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda"), + ) + + +def _hybrid_recipe_fp8_current_row_delayed_col(): + """Cross-format per-tensor Float8: current rowwise + delayed columnwise. + + Reversed variant of ``_hybrid_recipe_fp8_delayed_row_current_col``: row + sub-storage -> current bucket, col sub-storage -> delayed bucket. + """ + return _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E4M3), + grad_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E5M2), + ) + + +def _hybrid_recipe_mxfp8(): + """Same-format MXFP8 on both directions (rejected today; TODO #3158).""" + return _hybrid_custom_recipe( + row_factory=lambda: MXFP8Quantizer(tex.DType.kFloat8E4M3), + col_factory=lambda: MXFP8Quantizer(tex.DType.kFloat8E4M3), + grad_factory=lambda: MXFP8Quantizer(tex.DType.kFloat8E5M2), + ) + + +def _hybrid_recipe_blockwise(): + """Same-format Float8Blockwise on both directions (rejected today; TODO #3158).""" + return _hybrid_custom_recipe( + row_factory=lambda: Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True + ), + col_factory=lambda: Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True + ), + grad_factory=lambda: Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E5M2, rowwise=True, columnwise=True + ), + ) + + +def _build_hybrid_linear_weight(out_features, in_features, hybrid_recipe): + """Build a `HybridQuantizedTensor` weight via `quantized_model_init`. + + Returns (weight, fp32_high_precision_init_val) where the high-precision + init val is on GPU so we can use it as the "master" weight in + quantize_master_weights tests. + """ + torch.manual_seed(42) + with quantized_model_init( + enabled=True, + recipe=hybrid_recipe, + preserve_high_precision_init_val=True, + ): + model = Linear(in_features, out_features, bias=False, params_dtype=torch.bfloat16).cuda() + + weight = model.weight + assert isinstance( + weight, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(weight).__name__}" + hp_init_cpu = weight.get_high_precision_init_val() + assert hp_init_cpu is not None, "preserve_high_precision_init_val should populate the cpu val" + hp_init = hp_init_cpu.to(weight.device).float() + return weight, hp_init + + +def _hybrid_param_for(out_features, in_features, hybrid_recipe): + """Same as `_build_hybrid_linear_weight` but discards the init val.""" + weight, _ = _build_hybrid_linear_weight(out_features, in_features, hybrid_recipe) + return weight + + +@requires_fp8 +class TestHybridQuantizeMasterWeights: + """`quantize_master_weights` + `post_all_gather_processing` for hybrid params. + + Dispatch is per-direction: each sub-storage is routed independently into the + per-format bucket matching its own sub-quantizer type. Currently-supported + sub-quantizer types can mix freely across directions (e.g. Float8 delayed + row + Float8 current col), single-direction hybrid (one sub-storage dropped + via ``update_usage``) routes the live direction(s) only; per-block sub- + quantizers (MXFP8, NVFP4, Float8Blockwise) raise NotImplementedError + regardless of which direction they appear in. + + Supported subset (per-tensor Float8) -- positive tests verify the present + sub-storage(s) dequantize close to the master weight after the cast: + + * Float8CurrentScaling on both directions (same-format, full master) + * Float8 delayed scaling on both directions (same-format) + * Float8 delayed row + Float8 current col (cross-format; row -> delayed + bucket, col -> current bucket) + * Float8 current row + Float8 delayed col (cross-format, reversed) + * Single-direction (rowwise-only) hybrid via ``update_usage`` + * Single-direction (columnwise-only) hybrid via ``update_usage`` + + Rejected subset (NotImplementedError / ValueError) -- negative tests pin + the per-direction rejection contract and the both-None guardrail: + * MXFP8 as a hybrid sub-quantizer (rowwise OR columnwise) + * NVFP4 as a hybrid sub-quantizer (rowwise OR columnwise) + * Float8Blockwise as a hybrid sub-quantizer + * A live ``rowwise_dequantized`` column (deferred to #3158) + * A partial ``original``-source master without a two-direction FSDP shard + * Both sub-storages dropped (caller bug: nothing left to cast) + """ + + @staticmethod + def _make_transpose_only_float8_weight(shape, quantizer, *, fill_value=173): + """Build a Hopper/L40-style columnwise-only Float8Tensor. + + Blackwell keeps ``_data`` populated for columnwise-only Float8, so these + tests synthesize the older architecture layout directly: logical + ``[M, K]`` shape with the only live FP8 bytes in ``_transpose[K, M]``. + """ + rows, cols = shape + return Float8Tensor( + shape=shape, + dtype=torch.bfloat16, + data=None, + data_transpose=torch.full((cols, rows), fill_value, dtype=torch.uint8, device="cuda"), + fp8_scale_inv=torch.ones(1, dtype=torch.float32, device="cuda"), + fp8_dtype=tex.DType.kFloat8E4M3, + quantizer=quantizer, + requires_grad=False, + device="cuda", + ) + + @staticmethod + def _scatter_expected_logical_bytes(initial_transpose, fp8_bytes, logical_shape, start_offset): + rows, cols = logical_shape + expected = initial_transpose.clone() + expected_2d = expected.reshape(cols, rows) + remaining = fp8_bytes.numel() + logical_offset = start_offset + src_offset = 0 + while remaining > 0: + row = logical_offset // cols + col = logical_offset % cols + n = min(remaining, cols - col) + expected_2d[col : col + n, row].copy_(fp8_bytes[src_offset : src_offset + n]) + logical_offset += n + src_offset += n + remaining -= n + return expected + + @staticmethod + def _reference_fp8_bytes(master, scale, dtype=torch.bfloat16): + quantizer = Float8Quantizer( + scale=scale, + amax=torch.zeros(1, dtype=torch.float32, device="cuda"), + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + ) + raw = torch.empty((1, master.numel()), dtype=torch.uint8, device="cuda") + temp = quantizer.create_tensor_from_data(raw, dtype) + quantizer.update_quantized(master.reshape(1, -1), temp) + return temp._data.reshape(-1) + + @pytest.mark.parametrize("partial_master", (False, True)) + def test_rowwise_dequantized_master_update_raises_before_mutation( + self, + partial_master, + ): + """Defer row-to-column sequencing to the #3158 follow-up.""" + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + shape = (4, 8) + quantizer = HybridQuantizer( + rowwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + ), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + source = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + weight = quantizer(source) + state_before = { + "rowwise": _snapshot_storage_tensor_metadata( + weight._rowwise_storage, + clone=True, + ), + "columnwise": _snapshot_storage_tensor_metadata( + weight._columnwise_storage, + clone=True, + ), + } + start_offset = shape[-1] if partial_master else 0 + master = torch.randn(weight.numel(), dtype=torch.float32, device="cuda") + master = master[start_offset:].contiguous() + + with pytest.raises( + NotImplementedError, + match="rowwise update/all-gather.*#3158", + ): + quantize_master_weights( + [weight], + [master], + [start_offset], + group=group, + ) + + state_after = { + "rowwise": _snapshot_storage_tensor_metadata( + weight._rowwise_storage, + clone=True, + ), + "columnwise": _snapshot_storage_tensor_metadata( + weight._columnwise_storage, + clone=True, + ), + } + _assert_nested_state_exact(state_after, state_before) + + def test_fp8_current_original_partial_master_raises(self): + """Reject an independently sourced Hybrid column from a partial master shard. + + A one-payload distributed optimizer only communicates the rowwise payload. + Re-creating an independently quantized column from a partial master shard can + therefore produce a different value after checkpoint resume. Construct the + Hopper transpose-only column explicitly so this safety regression does not + depend on columnwise-only kernel support. + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + shape = (4, 8) + row_quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + rowwise=True, + columnwise=False, + ) + col_quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + rowwise=False, + columnwise=True, + ) + quantizer = HybridQuantizer( + rowwise_quantizer=row_quantizer, + columnwise_quantizer=col_quantizer, + columnwise_source="original", + ) + + source = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + rowwise = row_quantizer(source) + columnwise = self._make_transpose_only_float8_weight(shape, col_quantizer) + weight = HybridQuantizedTensor( + shape=shape, + dtype=source.dtype, + rowwise_storage=rowwise, + columnwise_storage=columnwise, + quantizer=quantizer, + requires_grad=False, + device="cuda", + ) + master_shard = source.reshape(-1)[shape[-1] :].float().contiguous() + + with pytest.raises( + ValueError, + match="partial master shard.*columnwise_source='original'.*full-master data", + ): + quantize_master_weights( + [weight], + [master_shard], + [shape[-1]], + group=None, + ) + + @pytest.mark.skipif( + is_non_tn_fp8_gemm_supported(), + reason="Hopper-only: Blackwell supports columnwise per-tensor FP8 FSDP updates", + ) + @pytest.mark.parametrize("scaling", ("current", "delayed")) + def test_per_tensor_fp8_fsdp_transpose_only_raises_before_mutation(self, scaling): + """Reject Hopper FSDP updates that cannot flatten columnwise-only storage.""" + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + shape = (4, 8) + shard_shape = (2, 8) + if scaling == "current": + row_quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + rowwise=True, + columnwise=False, + ) + col_quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + rowwise=False, + columnwise=True, + ) + else: + row_quantizer = _make_delayed_quantizer(rowwise=True, columnwise=False) + col_quantizer = _make_delayed_quantizer(rowwise=False, columnwise=True) + quantizer = HybridQuantizer( + rowwise_quantizer=row_quantizer, + columnwise_quantizer=col_quantizer, + ) + + source = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + shard_source = source[: shard_shape[0]].contiguous() + weight = HybridQuantizedTensor( + shape=shape, + dtype=source.dtype, + rowwise_storage=row_quantizer(source), + columnwise_storage=self._make_transpose_only_float8_weight(shape, col_quantizer), + quantizer=quantizer, + requires_grad=False, + device="cuda", + ) + shard = HybridQuantizedTensor( + shape=shard_shape, + dtype=source.dtype, + rowwise_storage=row_quantizer(shard_source), + columnwise_storage=self._make_transpose_only_float8_weight(shard_shape, col_quantizer), + quantizer=quantizer, + requires_grad=False, + device="cuda", + ) + weight_col = weight._columnwise_storage + shard_col = shard._columnwise_storage + weight_scale_before = weight_col._scale_inv.clone() + weight_transpose_before = weight_col._transpose.clone() + shard_scale_before = shard_col._scale_inv.clone() + shard_transpose_before = shard_col._transpose.clone() + + with pytest.raises( + NotImplementedError, + match="Columnwise-only per-tensor FP8 quantization is not implemented", + ): + quantize_master_weights( + [weight], + [shard_source.float()], + [0], + group=None, + fsdp_shard_model_weights=[shard], + ) + + assert weight_col._data is None + assert shard_col._data is None + torch.testing.assert_close(weight_col._scale_inv, weight_scale_before, rtol=0, atol=0) + torch.testing.assert_close(shard_col._scale_inv, shard_scale_before, rtol=0, atol=0) + assert torch.equal(weight_col._transpose, weight_transpose_before) + assert torch.equal(shard_col._transpose, shard_transpose_before) + + @staticmethod + def _logical_float8_bytes(storage): + """Return FP8 payload in the tensor's logical row-major order.""" + if storage._data is not None: + return storage._data.reshape(-1) + assert storage._transpose is not None + return storage._transpose.transpose(-2, -1).contiguous().reshape(-1) + + def test_fp8_current_transpose_only_nonzero_offset(self): + """Current-scaling distopt update handles Hopper-style columnwise storage. + + Regression for ``model_weight.reshape(-1)`` reaching + ``Float8Tensor._ReshapeFunc.forward`` and dereferencing ``_data=None``. + The nonzero offset spans multiple logical rows, so the test also checks + row-major shard bytes are scattered into transposed storage correctly. + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + shape = (4, 8) + start_offset = 5 + master = torch.linspace(-2.0, 2.0, steps=17, dtype=torch.float32, device="cuda") + quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda", rowwise=False, columnwise=True + ) + weight = self._make_transpose_only_float8_weight(shape, quantizer) + initial = weight._transpose.clone() + + quantize_master_weights([weight], [master], [start_offset], group=group) + + scale = torch.reciprocal(weight._scale_inv.detach().clone()) + fp8_bytes = self._reference_fp8_bytes(master.to(weight.dtype), scale, weight.dtype) + expected = self._scatter_expected_logical_bytes(initial, fp8_bytes, shape, start_offset) + assert weight._data is None + assert weight._transpose_invalid is False + assert torch.equal(weight._transpose, expected) + + def test_fp8_delayed_transpose_only_nonzero_offset(self): + """Delayed-scaling distopt update handles Hopper-style columnwise storage. + + Regression for the direct ``model_weight._data.view(-1)`` path in the + delayed-scaling helper. + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + shape = (4, 8) + start_offset = 6 + master = torch.linspace(-3.0, 1.0, steps=15, dtype=torch.float32, device="cuda") + quantizer = _make_delayed_quantizer(tex.DType.kFloat8E4M3) + quantizer.set_usage(rowwise=False, columnwise=True) + weight = self._make_transpose_only_float8_weight(shape, quantizer) + initial = weight._transpose.clone() + + quantize_master_weights([weight], [master], [start_offset], group=group) + + fp8_bytes = self._reference_fp8_bytes( + master.to(weight.dtype), weight._get_quantizer().scale, weight.dtype + ) + expected = self._scatter_expected_logical_bytes(initial, fp8_bytes, shape, start_offset) + assert weight._data is None + assert weight._transpose_invalid is False + assert torch.equal(weight._transpose, expected) + + # ---------- Positive tests (same-format) ---------- + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_fp8_current_same_format_full_master(self): + """Full master (start_offset=0) routes both sub-storages through the + existing per-format current-scaling helper. Verifies both directions + dequantize close to the master weight after the cast. + """ + from transformer_engine.pytorch.tensor.utils import ( + quantize_master_weights, + post_all_gather_processing, + ) + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_current() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + # Distributed-optimizer convention: master weight is the flat FP32 shard + # owned by the current rank (or the full param for non-distributed cases). + master_flat = hp_master.view(-1).contiguous() + + quantize_master_weights([weight], [master_flat], [0], group=group) + post_all_gather_processing([weight]) + + assert weight._rowwise_storage is not None + assert weight._columnwise_storage is not None + master_bf16 = master_flat.to(weight.dtype).reshape(weight.shape) + expected_row = weight._quantizer.rowwise_quantizer.copy().quantize(master_bf16) + expected_column = weight._quantizer.columnwise_quantizer.copy().quantize(master_bf16) + _assert_storage_data_exact( + weight._rowwise_storage, + expected_row, + context="full-master rowwise independent quantization", + ) + _assert_storage_data_exact( + weight._columnwise_storage, + expected_column, + context="full-master columnwise independent quantization", + ) + dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) + dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) + # FP8 E4M3 round-trip; matches the loose tolerance the equivalent + # native-FP8-current test uses (e.g. test_dequantize_close_to_original). + torch.testing.assert_close(dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1) + torch.testing.assert_close(dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1) + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_fp8_delayed_same_format_full_master(self): + """Same-format delayed scaling on both directions. Both sub-storages + route into the delayed-scaling bucket as independent entries; the + helper processes them with a single bucket-wide amax all-reduce. + Verifies each direction dequantizes close to the master weight. + """ + from transformer_engine.pytorch.tensor.utils import ( + quantize_master_weights, + post_all_gather_processing, + ) + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_delayed() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + quantize_master_weights([weight], [master_flat], [0], group=group) + post_all_gather_processing([weight]) + + assert weight._rowwise_storage is not None + assert weight._columnwise_storage is not None + dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) + dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) + torch.testing.assert_close(dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1) + torch.testing.assert_close(dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1) + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_fp8_delayed_row_current_col_full_master(self): + """Cross-format per-tensor Float8: delayed row + current col. + + Pins the new per-direction routing: row sub-storage goes to the + delayed bucket, col sub-storage goes to the current bucket. Each + helper runs independently on its single-entry bucket, with no + cross-pollination between the two scaling lifecycles. + """ + from transformer_engine.pytorch.tensor.utils import ( + quantize_master_weights, + post_all_gather_processing, + ) + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_delayed_row_current_col() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + quantize_master_weights([weight], [master_flat], [0], group=group) + post_all_gather_processing([weight]) + + assert weight._rowwise_storage is not None + assert weight._columnwise_storage is not None + assert isinstance(weight._quantizer.rowwise_quantizer, Float8Quantizer) + assert isinstance(weight._quantizer.columnwise_quantizer, Float8CurrentScalingQuantizer) + dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) + dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) + torch.testing.assert_close(dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1) + torch.testing.assert_close(dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1) + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_fp8_current_row_delayed_col_full_master(self): + """Cross-format per-tensor Float8: current row + delayed col. + + Reversed variant of the test above — pins that the per-direction + loop's second iteration (col) reaches the delayed dispatch arm + independently of what the rowwise iteration did. + """ + from transformer_engine.pytorch.tensor.utils import ( + quantize_master_weights, + post_all_gather_processing, + ) + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_current_row_delayed_col() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + quantize_master_weights([weight], [master_flat], [0], group=group) + post_all_gather_processing([weight]) + + assert weight._rowwise_storage is not None + assert weight._columnwise_storage is not None + assert isinstance(weight._quantizer.rowwise_quantizer, Float8CurrentScalingQuantizer) + assert isinstance(weight._quantizer.columnwise_quantizer, Float8Quantizer) + dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) + dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) + torch.testing.assert_close(dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1) + torch.testing.assert_close(dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1) + + # NOTE: Per-block sub-quantizers (MXFP8, NVFP4, Float8Blockwise) are not + # supported as hybrid sub-quantizers by this initial integration, regardless + # of which direction they appear in. See the per-direction rejection tests + # below (``test_mxfp8_*_raises`` covers both rowwise and columnwise rejection + # of MXFP8; ``test_nvfp4_*_raises`` and ``test_blockwise_*_raises`` similarly). + # The TODO #3158 block above ``_route_hybrid_to_buckets`` in tensor/utils.py + # documents the upstream constraints (single-direction cast helper / kernel + # support) whose unblocker drops per-block format support in for free. + + # ---------- Negative tests (per-direction rejection contract) ---------- + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_mxfp8_rowwise_raises(self): + """MXFP8 in the rowwise sub-quantizer is rejected per-direction. + + ``_cast_master_weights_to_fp8_mxfp8_scaling`` assumes each entry's + ``model_weight`` has BOTH ``_rowwise_*`` and ``_columnwise_*`` populated + (the underlying partial-cast kernel is bidirectional), while a hybrid + sub-storage is single-direction by construction. See + ``TODO(#3158, hybrid-mxfp8-distopt)`` in tensor/utils.py for the unblocker shape. + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_mxfp8() + # Shape must be a multiple of MXFP8 block size (32) on both axes. + weight, hp_master = _build_hybrid_linear_weight(64, 128, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + with pytest.raises(NotImplementedError, match="MXFP8Quantizer rowwise"): + quantize_master_weights([weight], [master_flat], [0], group=group) + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_mxfp8_columnwise_raises(self): + """MXFP8 in the columnwise sub-quantizer is rejected per-direction. + + Pairs FP8 current scaling in the rowwise slot (supported) with MXFP8 + in the columnwise slot (rejected). The rowwise iteration of + ``_route_hybrid_to_buckets`` routes the FP8 sub-storage into the + current-scaling bucket cleanly; the columnwise iteration then hits + MXFP8 and raises. Pins that per-direction dispatch visits and rejects + the columnwise sub-quantizer too — not just the rowwise one. + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: MXFP8Quantizer(tex.DType.kFloat8E4M3), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + # Shape must be a multiple of MXFP8 block size (32) on both axes. + weight, hp_master = _build_hybrid_linear_weight(64, 128, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + with pytest.raises(NotImplementedError, match="MXFP8Quantizer columnwise"): + quantize_master_weights([weight], [master_flat], [0], group=group) + + @pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}") + def test_nvfp4_rowwise_raises(self): + """NVFP4 in the rowwise sub-quantizer is rejected per-direction. + + The NVFP4 cast path is blocked on a pair of upstream constraints + documented in the TODO #3158 block above ``_route_hybrid_to_buckets`` in + tensor/utils.py. + + NOTE: after PR #3027, single-direction 2D NVFP4 construction works, + so this test now reaches the intended ``quantize_master_weights`` + rejection while using the base weight scaling mode. + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, with_2d_quantization=True + ), + col_factory=lambda: NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, with_2d_quantization=True + ), + grad_factory=lambda: NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, with_2d_quantization=False + ), + ) + weight, hp_master = _build_hybrid_linear_weight(64, 128, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + with pytest.raises(NotImplementedError, match="NVFP4Quantizer rowwise"): + quantize_master_weights([weight], [master_flat], [0], group=group) + + @pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}") + def test_nvfp4_columnwise_raises(self): + """NVFP4 in only the columnwise slot is rejected per-direction.""" + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, with_2d_quantization=True + ), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + weight, hp_master = _build_hybrid_linear_weight(64, 128, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + with pytest.raises(NotImplementedError, match="NVFP4Quantizer columnwise"): + quantize_master_weights([weight], [master_flat], [0], group=group) + + @pytest.mark.skipif( + not fp8_block_scaling_available, + reason=f"Float8 block scaling: {reason_for_no_fp8_block_scaling}", + ) + def test_blockwise_rowwise_raises(self): + """Float8BlockQuantizer in the rowwise sub-quantizer is rejected + per-direction (no e2e factory uses it; TODO #3158 marker in tensor/utils.py). + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_blockwise() + weight, hp_master = _build_hybrid_linear_weight(128, 128, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + with pytest.raises(NotImplementedError, match="Float8BlockQuantizer rowwise"): + quantize_master_weights([weight], [master_flat], [0], group=group) + + @pytest.mark.skipif( + not fp8_block_scaling_available, + reason=f"Float8 block scaling: {reason_for_no_fp8_block_scaling}", + ) + def test_blockwise_columnwise_raises(self): + """Float8BlockQuantizer in only the columnwise slot is rejected.""" + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=2, + ), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + weight, hp_master = _build_hybrid_linear_weight(128, 128, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + with pytest.raises(NotImplementedError, match="Float8BlockQuantizer columnwise"): + quantize_master_weights([weight], [master_flat], [0], group=group) + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_rowwise_only_fp8_current_full_master(self): + """Single-direction hybrid: columnwise dropped via update_usage. + + Pins that the per-direction loop in `_route_hybrid_to_buckets` skips + the dropped direction silently and routes only the present (rowwise) + sub-storage. Useful for inference / memory-saving paths that + deliberately keep only the fprop-side direction. + """ + from transformer_engine.pytorch.tensor.utils import ( + quantize_master_weights, + post_all_gather_processing, + ) + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_current() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + weight.update_usage(rowwise_usage=True, columnwise_usage=False) + assert weight._rowwise_storage is not None + assert weight._columnwise_storage is None + master_flat = hp_master.view(-1).contiguous() + + quantize_master_weights([weight], [master_flat], [0], group=group) + post_all_gather_processing([weight]) + + # Columnwise stays dropped (the cast must not silently revive it). + assert weight._columnwise_storage is None + expected_row = weight._quantizer.rowwise_quantizer.copy().quantize( + master_flat.to(weight.dtype).reshape(weight.shape) + ) + _assert_storage_data_exact( + weight._rowwise_storage, + expected_row, + context="rowwise-only full-master independent quantization", + ) + # Rowwise is populated and dequantizes close to the master. + dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) + torch.testing.assert_close(dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1) + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_columnwise_only_fp8_current_full_master(self): + """Single-direction hybrid: rowwise dropped via update_usage. + + Reversed variant — verifies the column-only iteration of the per- + direction loop reaches the dispatch and routes correctly. + """ + from transformer_engine.pytorch.tensor.utils import ( + quantize_master_weights, + post_all_gather_processing, + ) + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_current() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + weight.update_usage(rowwise_usage=False, columnwise_usage=True) + assert weight._rowwise_storage is None + assert weight._columnwise_storage is not None + master_flat = hp_master.view(-1).contiguous() + + quantize_master_weights([weight], [master_flat], [0], group=group) + post_all_gather_processing([weight]) + + # Rowwise stays dropped (the cast must not silently revive it). + assert weight._rowwise_storage is None + expected_column = weight._quantizer.columnwise_quantizer.copy().quantize( + master_flat.to(weight.dtype).reshape(weight.shape) + ) + _assert_storage_data_exact( + weight._columnwise_storage, + expected_column, + context="columnwise-only full-master independent quantization", + ) + # Columnwise is populated and dequantizes close to the master. + dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) + torch.testing.assert_close(dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1) + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_both_sub_storages_none_raises(self): + """Both sub-storages dropped via update_usage — nothing left to cast. + + This is the only remaining sub-storage-presence guardrail after the + single-direction enablement: a fully-dropped hybrid weight reaching + `quantize_master_weights` is a caller bug, not a deferred feature, + so we surface it as a ValueError. + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_current() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + weight.update_usage(rowwise_usage=False, columnwise_usage=False) + assert weight._rowwise_storage is None + assert weight._columnwise_storage is None + master_flat = hp_master.view(-1).contiguous() + + with pytest.raises(ValueError, match="both rowwise and columnwise"): + quantize_master_weights([weight], [master_flat], [0], group=group) + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridPostAllGatherProcessing: + """Hybrid branch of `post_all_gather_processing` is exercised indirectly by + the positive `TestHybridQuantizeMasterWeights` tests; the case below pins + an additional invariant that the routing logic must preserve. + """ + + def test_post_ag_idempotent_for_fp8_current_hybrid(self): + """Calling post_all_gather_processing twice on a same-format Float8 + hybrid must not corrupt the sub-storages. + """ + from transformer_engine.pytorch.tensor.utils import ( + quantize_master_weights, + post_all_gather_processing, + ) + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_current() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + quantize_master_weights([weight], [master_flat], [0], group=group) + post_all_gather_processing([weight]) + dq_row_first = weight._rowwise_storage.dequantize(dtype=torch.float32) + dq_col_first = weight._columnwise_storage.dequantize(dtype=torch.float32) + + post_all_gather_processing([weight]) + dq_row_second = weight._rowwise_storage.dequantize(dtype=torch.float32) + dq_col_second = weight._columnwise_storage.dequantize(dtype=torch.float32) + + torch.testing.assert_close(dq_row_first, dq_row_second, rtol=0.0, atol=0.0) + torch.testing.assert_close(dq_col_first, dq_col_second, rtol=0.0, atol=0.0) + + +# --------------------------------------------------------------------------- +# 4. Recipe correspondence validation +# --------------------------------------------------------------------------- + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridRecipeCorrespondence: + """Test _check_weight_tensor_recipe_correspondence with hybrid params.""" + + def _hybrid_fp8_recipe(self): + return _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + + def test_hybrid_param_with_matching_recipe_does_not_raise(self): + """Forward pass with matching recipe should not raise.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + out = model(inp) + assert not torch.isnan(out).any() + + def test_hybrid_param_with_mismatched_recipe_raises(self): + """Forward pass with a non-CustomRecipe on a hybrid param should raise.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + mismatch_recipe = recipe.Float8CurrentScaling() + with pytest.raises(RuntimeError, match="Recipe mismatch"): + with torch.no_grad(): + with autocast(enabled=True, recipe=mismatch_recipe): + model(inp) + + +# --------------------------------------------------------------------------- +# 5. quantize_ in-place update for hybrid +# --------------------------------------------------------------------------- + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridQuantizeInPlace: + """Test in-place re-quantization (quantize_) for HybridQuantizedTensor. + + This is needed for the optimizer writeback path (param.quantize_(master_weight)) + and the workspace cache update path (out.quantize_(new_bf16_weight)). + """ + + def test_quantize_inplace_updates_data(self): + """quantize_() should re-quantize both sub-storages from new BF16 data.""" + torch.manual_seed(42) + hq = HybridQuantizer( + rowwise_quantizer=Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + columnwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + ) + original = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + tensor = hq.quantize(original) + + # Update with different data + new_data = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + expected = hq.quantize(new_data) + result = tensor.quantize_(new_data) + + assert result is tensor + _assert_hybrid_tensor_exact(tensor, expected, context="quantize_") + + def test_quantize_inplace_preserves_tensor_identity(self): + """quantize_() should update in-place, not create a new tensor.""" + hq = HybridQuantizer( + rowwise_quantizer=Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + columnwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + ) + original = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + tensor = hq.quantize(original) + + new_data = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + result = tensor.quantize_(new_data) + + assert result is tensor, "quantize_() must return the object it updated" + + # noop_flag is a delayed-scaling feature; not tested here since + # delayed scaling is out of scope for hybrid quantization. + + +# --------------------------------------------------------------------------- +# 6. FusedAdam with hybrid quantized params +# --------------------------------------------------------------------------- + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridFusedAdam: + """Test FusedAdam optimizer with HybridQuantizedTensor parameters.""" + + def _build_hybrid_model(self): + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + return model, hybrid_recipe + + def test_fused_adam_accepts_hybrid_params(self): + """FusedAdam should not crash when given HybridQuantizedTensor params.""" + model, _ = self._build_hybrid_model() + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + assert optimizer is not None + + def test_fused_adam_master_weights_track_reference(self): + """FP32 master weights should closely track a reference Adam optimizer. + + Small divergence is expected because HybridQuantizedTensor.float() + may take a slightly different dequantization path than + detach().clone().float() through __torch_dispatch__. + """ + model, _ = self._build_hybrid_model() + + ref_params = [p.detach().clone().float() for p in model.parameters()] + + options = {"lr": 5e-4, "betas": (0.9, 0.999), "eps": 1e-8, "weight_decay": 0} + ref_optim = torch.optim.Adam(ref_params, **options) + tst_optim = te.optimizers.FusedAdam( + list(model.parameters()), + master_weights=True, + master_weight_dtype=torch.float32, + use_decoupled_grad=True, + **options, + ) + + for _ in range(5): + for p_ref, p in zip(ref_params, model.parameters()): + p_ref.grad = torch.rand_like(p_ref) + p.decoupled_grad = p_ref.grad.clone() + ref_optim.step() + tst_optim.step() + + master_params = [ + tst_optim.get_unscaled_state(p, "master_param") for p in model.parameters() + ] + torch.testing.assert_close(ref_params, master_params, rtol=1e-3, atol=1e-3) + + def test_fused_adam_param_remains_hybrid_after_step(self): + """Weight params should still be HybridQuantizedTensors after optimizer step.""" + model, _ = self._build_hybrid_model() + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + use_decoupled_grad=True, + ) + + for _ in range(3): + for p in model.parameters(): + p.decoupled_grad = torch.rand_like(p.float()) + optimizer.step() + + for name, p in model.named_parameters(): + if "bias" not in name: + assert isinstance( + p, HybridQuantizedTensor + ), f"{name} lost HybridQuantizedTensor type: {type(p).__name__}" + + def test_fused_adam_requires_master_weights(self): + """FusedAdam without master_weights should raise for hybrid quantized params.""" + model, _ = self._build_hybrid_model() + + with pytest.raises(RuntimeError, match="master_weights"): + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=False, + ) + for p in model.parameters(): + p.grad = torch.rand_like(p.float()).to(p.dtype) + optimizer.step() + + +# --------------------------------------------------------------------------- +# 7. End-to-end training loop: fwd + bwd + optimizer step +# --------------------------------------------------------------------------- + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridQuantizedParamsEndToEnd: + """Full training loop: quantized_model_init + autocast fwd + bwd + FusedAdam.step().""" + + def _build_model_and_recipe(self): + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + return model, hybrid_recipe + + def test_training_loop_loss_decreases(self): + """Loss should decrease over a few training steps.""" + torch.manual_seed(42) + model, hybrid_recipe = self._build_model_and_recipe() + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(4, 32, 256, dtype=torch.bfloat16, device="cuda") + target = torch.randn_like(x) + + losses = [] + for i in range(7): + optimizer.zero_grad(set_to_none=True) + with autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + loss = torch.nn.functional.mse_loss(output, target) + losses.append(loss.item()) + loss.backward() + + for name, p in model.named_parameters(): + assert p.grad is not None, f"Step {i}: {name} has no gradient" + assert torch.isfinite(p.grad).all(), f"Step {i}: {name} has non-finite grad" + + optimizer.step() + + # Strictly monotonic decrease + assert all( + losses[i + 1] < losses[i] for i in range(len(losses) - 1) + ), f"Loss not strictly decreasing each step: {losses}" + + def test_training_loop_params_remain_quantized(self): + """Params should remain HybridQuantizedTensors after training.""" + torch.manual_seed(42) + model, hybrid_recipe = self._build_model_and_recipe() + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(4, 32, 256, dtype=torch.bfloat16, device="cuda") + target = torch.randn_like(x) + + for _ in range(3): + optimizer.zero_grad(set_to_none=True) + with autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + loss = torch.nn.functional.mse_loss(output, target) + loss.backward() + optimizer.step() + + for name, p in model.named_parameters(): + if "bias" not in name: + assert isinstance( + p, HybridQuantizedTensor + ), f"{name} is {type(p).__name__}, not HybridQuantizedTensor" + + def test_training_loop_optimizer_states_are_fp32(self): + """Optimizer states should be FP32.""" + torch.manual_seed(42) + model, hybrid_recipe = self._build_model_and_recipe() + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(4, 32, 256, dtype=torch.bfloat16, device="cuda") + for _ in range(2): + optimizer.zero_grad(set_to_none=True) + with autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + output.float().sum().backward() + optimizer.step() + + for name, p in model.named_parameters(): + state = optimizer.state[p] + assert state["exp_avg"].dtype == torch.float32 + assert state["exp_avg_sq"].dtype == torch.float32 + if "bias" not in name: + assert state["master_param"].dtype == torch.float32 + + +# --------------------------------------------------------------------------- +# 8. Mixed-format quantized params (e.g. MXFP8 row + NVFP4 col) +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif( + not (mxfp8_available and nvfp4_available), + reason=f"MXFP8: {reason_for_no_mxfp8}; NVFP4: {reason_for_no_nvfp4}", +) +class TestHybridMixedFormatQuantizedParams: + """Quantized params with genuinely different formats per direction.""" + + def _build_mixed_model(self, in_features=256, out_features=256): + """MXFP8 rowwise (fprop) + role-aware NVFP4 columnwise (bwd).""" + + def qfactory(role): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight", "output"): + return HybridQuantizer( + rowwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + columnwise_quantizer=nvfp4_factory(role), + ) + if is_linear and role.tensor_type == "grad_output": + return nvfp4_factory(role) + return MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + + hybrid_recipe = recipe.CustomRecipe(qfactory=qfactory) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + return model, hybrid_recipe + + def test_mixed_format_param_creation(self): + """Model init with mixed MXFP8/NVFP4 hybrid should produce a + HybridQuantizedTensor parameter.""" + model, _ = self._build_mixed_model() + assert isinstance(model.weight, HybridQuantizedTensor) + + def test_mixed_format_forward_only(self): + """Forward pass with mixed-format quantized params.""" + torch.manual_seed(42) + model, hybrid_recipe = self._build_mixed_model() + inp = torch.randn(32, 256, device="cuda", dtype=torch.bfloat16) + + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + out = model(inp) + + assert out.shape == (32, 256) + assert not torch.isnan(out).any() + assert not torch.isinf(out).any() + + def test_mixed_format_forward_backward(self): + """Full fwd+bwd with mixed-format quantized params.""" + torch.manual_seed(42) + model, hybrid_recipe = self._build_mixed_model() + inp = torch.randn(32, 256, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + with autocast(enabled=True, recipe=hybrid_recipe): + out = model(inp) + loss = out.float().sum() + loss.backward() + + assert inp.grad is not None + assert not torch.isnan(inp.grad).any() + for name, p in model.named_parameters(): + assert p.grad is not None, f"No gradient for {name}" + assert not torch.isnan(p.grad).any(), f"NaN gradient for {name}" + + def test_mixed_format_training_loop(self): + """End-to-end training loop with mixed-format hybrid quantized params.""" + torch.manual_seed(42) + model, hybrid_recipe = self._build_mixed_model() + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(4, 32, 256, dtype=torch.bfloat16, device="cuda") + target = torch.randn_like(x) + + losses = [] + for i in range(5): + optimizer.zero_grad(set_to_none=True) + with autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + loss = torch.nn.functional.mse_loss(output, target) + losses.append(loss.item()) + loss.backward() + optimizer.step() + + # Strictly monotonic decrease + assert all( + losses[i + 1] < losses[i] for i in range(len(losses) - 1) + ), f"Loss not strictly decreasing each step: {losses}" + for name, p in model.named_parameters(): + if "bias" not in name: + assert isinstance(p, HybridQuantizedTensor), f"{name} is {type(p).__name__}" + + def test_mixed_format_sub_storage_types(self): + """Verify that sub-storages have the correct types (MXFP8 vs NVFP4).""" + model, _ = self._build_mixed_model() + weight = model.weight + from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import ( + MXFP8TensorStorage, + ) + + row = weight.rowwise_sub_storage + col = weight.columnwise_sub_storage + assert isinstance(row, MXFP8TensorStorage) or hasattr( + row, "_rowwise_data" + ), f"Expected MXFP8 rowwise sub-storage, got {type(row).__name__}" + assert isinstance(col, NVFP4TensorStorage) or hasattr( + col, "_rowwise_data" + ), f"Expected NVFP4 columnwise sub-storage, got {type(col).__name__}" + + +# --------------------------------------------------------------------------- +# 9. Quantized params equivalence: vanilla vs hybrid (same format both dirs) +# --------------------------------------------------------------------------- + + +def _hybrid_mxfp8_qfactory(role): + """Hybrid MXFP8 (E4M3 both dirs).""" + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + return HybridQuantizer( + rowwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + columnwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + ) + + +def _hybrid_nvfp4_qfactory(role): + """Hybrid NVFP4 (E2M1 both dirs, base role behavior).""" + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type == "grad_output": + return nvfp4_factory(role) + return HybridQuantizer( + rowwise_quantizer=nvfp4_factory(role), + columnwise_quantizer=nvfp4_factory(role), + ) + + +class _QuantizedParamsEquivalenceBase: + """Base for comparing vanilla vs hybrid quantized params training. + + When the hybrid quantizer uses the same format in both directions, + the full quantized_model_init + training loop should produce + equivalent results to the vanilla (non-hybrid) quantized params path. + """ + + hidden_size = 256 + num_steps = 5 + + def _vanilla_recipe(self): + raise NotImplementedError + + def _hybrid_recipe(self): + raise NotImplementedError + + def _build_models(self): + """Create two models with identical init: one vanilla, one hybrid.""" + torch.manual_seed(42) + with quantized_model_init(enabled=True, recipe=self._vanilla_recipe()): + model_ref = Linear( + self.hidden_size, + self.hidden_size, + params_dtype=torch.bfloat16, + ).cuda() + + torch.manual_seed(42) + with quantized_model_init(enabled=True, recipe=self._hybrid_recipe()): + model_hyb = Linear( + self.hidden_size, + self.hidden_size, + params_dtype=torch.bfloat16, + ).cuda() + + return model_ref, model_hyb + + def _run_training_loop(self, model, train_recipe, x, target, num_steps): + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + trajectory = [] + hybrid_metadata_trajectory = [] + for step in range(num_steps): + optimizer.zero_grad(set_to_none=True) + step_input = x.detach().clone().requires_grad_(True) + _set_quantization_test_seed(199 + step) + with autocast(enabled=True, recipe=train_recipe): + output = model(step_input) + loss = torch.nn.functional.mse_loss(output, target) + loss.backward() + gradients = { + name: param.grad.detach().clone() + for name, param in model.named_parameters() + if param.grad is not None + } + pre_step_cuda_rng_state = torch.cuda.get_rng_state() + optimizer.step() + post_step_cuda_rng_state = torch.cuda.get_rng_state() + + hybrid_parameters = [ + (name, param) + for name, param in model.named_parameters() + if isinstance(param, HybridQuantizedTensor) + ] + expected_storages = {} + torch.cuda.set_rng_state(pre_step_cuda_rng_state) + try: + for name, param in hybrid_parameters: + master = optimizer.get_unscaled_state(param, "master_param") + expected_storages[name] = ( + param._quantizer.rowwise_quantizer.copy().quantize(master), + param._quantizer.columnwise_quantizer.copy().quantize(master), + ) + finally: + # The oracle must be observational only. Restore the state left + # by the real optimizer step after replaying NVFP4 stochastic + # rounding from its pre-step RNG state. + torch.cuda.set_rng_state(post_step_cuda_rng_state) + + hybrid_metadata = {} + for name, param in hybrid_parameters: + expected_row, expected_column = expected_storages[name] + _assert_storage_data_exact( + param.rowwise_sub_storage, + expected_row, + context=f"step {step} {name} rowwise writeback", + ) + _assert_storage_data_exact( + param.columnwise_sub_storage, + expected_column, + context=f"step {step} {name} columnwise writeback", + ) + hybrid_metadata[name] = { + "rowwise": _snapshot_storage_tensor_metadata( + param.rowwise_sub_storage, clone=True + ), + "columnwise": _snapshot_storage_tensor_metadata( + param.columnwise_sub_storage, clone=True + ), + } + hybrid_metadata_trajectory.append(hybrid_metadata) + logical_parameters = { + name: ( + param.dequantize(dtype=torch.float32).detach().clone() + if isinstance(param, QuantizedTensor) + else param.detach().float().clone() + ) + for name, param in model.named_parameters() + } + trajectory.append( + { + "output": output.detach().clone(), + "loss": loss.detach().clone(), + "input_gradient": step_input.grad.detach().clone(), + "parameter_gradients": gradients, + "logical_parameters": logical_parameters, + "optimizer": _clone_nested_state(optimizer.state_dict()), + } + ) + return trajectory, hybrid_metadata_trajectory + + def _test_equivalence(self): + model_ref, model_hyb = self._build_models() + + torch.manual_seed(99) + x = torch.randn(4, 32, self.hidden_size, dtype=torch.bfloat16, device="cuda") + target = torch.randn_like(x) + + ref_trajectory, ref_hybrid_metadata = self._run_training_loop( + model_ref, + self._vanilla_recipe(), + x, + target, + self.num_steps, + ) + hybrid_trajectory, hybrid_metadata_trajectory = self._run_training_loop( + model_hyb, + self._hybrid_recipe(), + x, + target, + self.num_steps, + ) + _assert_nested_state_exact( + hybrid_trajectory, + ref_trajectory, + path="same-format quantized-parameter trajectory", + ) + assert all(not step_metadata for step_metadata in ref_hybrid_metadata) + assert len(hybrid_metadata_trajectory) == self.num_steps + for step_metadata in hybrid_metadata_trajectory: + assert step_metadata.keys() == {"weight"} + assert step_metadata["weight"]["rowwise"] is not None + assert step_metadata["weight"]["columnwise"] is not None + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestQuantizedParamsEquivalenceFP8CurrentScaling(_QuantizedParamsEquivalenceBase): + """Vanilla versus same-format hybrid FP8-current training parity.""" + + def _vanilla_recipe(self): + return recipe.Float8CurrentScaling() + + def _hybrid_recipe(self): + return recipe.CustomRecipe(qfactory=_hybrid_fp8_current_qfactory) + + def test_equivalence(self): + self._test_equivalence() + + +@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") +class TestQuantizedParamsEquivalenceMXFP8(_QuantizedParamsEquivalenceBase): + """Vanilla MXFP8BlockScaling vs hybrid MXFP8 (same format both dirs).""" + + def _vanilla_recipe(self): + return recipe.MXFP8BlockScaling() + + def _hybrid_recipe(self): + return recipe.CustomRecipe(qfactory=_hybrid_mxfp8_qfactory) + + def test_equivalence(self): + self._test_equivalence() + + +@pytest.mark.skipif(not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling) +class TestQuantizedParamsEquivalenceBlockFP8(_QuantizedParamsEquivalenceBase): + """Vanilla Float8BlockScaling vs hybrid block FP8 (same format both dirs).""" + + def _vanilla_recipe(self): + return recipe.Float8BlockScaling() + + def _hybrid_recipe(self): + return recipe.CustomRecipe(qfactory=_hybrid_block_fp8_qfactory) + + def test_equivalence(self): + self._test_equivalence() + + +@pytest.mark.skipif( + not (fp8_available and nvfp4_available), + reason=f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", +) +class TestQuantizedParamsEquivalenceNVFP4(_QuantizedParamsEquivalenceBase): + """Vanilla NVFP4BlockScaling vs same-format hybrid NVFP4.""" + + def _vanilla_recipe(self): + return recipe.NVFP4BlockScaling() + + def _hybrid_recipe(self): + return recipe.CustomRecipe(qfactory=_hybrid_nvfp4_qfactory) + + def test_equivalence(self): + self._test_equivalence() + + +# --------------------------------------------------------------------------- +# 10. State dict save/load (checkpointing) for hybrid quantized params +# --------------------------------------------------------------------------- + + +# Module-level qfactories give TE-to-TE quantized-param checkpoints a stable +# importable reference for any pickled quantizer/recipe metadata. Portable BF16 +# checkpoint loading should not depend on importing these factories. + + +@requires_mxfp8_and_nvfp4 +class TestHybridCheckpoint: + """Test state_dict save/load round-trips for models with hybrid quantized params.""" + + def _hybrid_checkpoint_recipe(self): + return recipe.CustomRecipe(qfactory=mxfp8_fwd_nvfp4_bwd_factory) + + def test_state_dict_save_load_roundtrip(self): + """state_dict → save → load → same model should produce identical outputs.""" + torch.manual_seed(42) + hybrid_recipe = self._hybrid_checkpoint_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + out_before = model(inp) + + state_dict = model.state_dict() + + # Create a fresh model and load + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model2 = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + model2.load_state_dict(state_dict) + + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + out_after = model2(inp) + + torch.testing.assert_close(out_before, out_after, rtol=0.0, atol=0.0) + + def test_state_dict_contains_weight(self): + """state_dict should contain the weight key.""" + hybrid_recipe = self._hybrid_checkpoint_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + sd = model.state_dict() + assert "weight" in sd, f"state_dict keys: {list(sd.keys())}" + + def test_load_bf16_state_dict_into_hybrid_model(self): + """Loading a BF16 state_dict into a hybrid quantized model should work. + + This is the common scenario: pretrained BF16 weights loaded into a + model initialized with quantized_model_init. + """ + torch.manual_seed(42) + hybrid_recipe = self._hybrid_checkpoint_recipe() + + # Create BF16 model and get its state_dict + ref_model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + bf16_state_dict = ref_model.state_dict() + + # Create hybrid quantized model + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + # Load BF16 weights into hybrid model + model.load_state_dict(bf16_state_dict) + + # Verify model produces valid output + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + out = model(inp) + assert not torch.isnan(out).any() + assert not torch.isinf(out).any() + + def test_state_dict_torch_save_load(self): + """Full round-trip through torch.save/torch.load (file-based).""" + import tempfile + import os + + torch.manual_seed(42) + hybrid_recipe = self._hybrid_checkpoint_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + out_before = model(inp) + + with tempfile.NamedTemporaryFile(delete=False, suffix=".pt") as f: + torch.save(model.state_dict(), f.name) + tmp_path = f.name + + try: + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model2 = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + state_dict = torch.load(tmp_path, weights_only=False) + model2.load_state_dict(state_dict) + + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + out_after = model2(inp) + + torch.testing.assert_close(out_before, out_after, rtol=0.0, atol=0.0) + finally: + os.unlink(tmp_path) + + @staticmethod + def _checkpoint_training_step(model, optimizer, x, target, hybrid_recipe): + optimizer.zero_grad(set_to_none=True) + step_input = x.detach().clone().requires_grad_(True) + with autocast(enabled=True, recipe=hybrid_recipe): + output = model(step_input) + loss = torch.nn.functional.mse_loss(output, target) + loss.backward() + gradients = { + name: param.grad.detach().clone() + for name, param in model.named_parameters() + if param.grad is not None + } + optimizer.step() + return { + "output": output.detach().clone(), + "loss": loss.detach().clone(), + "input_gradient": step_input.grad.detach().clone(), + "gradients": gradients, + "parameters": _snapshot_model_parameters(model), + "optimizer": _clone_nested_state(optimizer.state_dict()), + } + + def test_checkpoint_resume_training(self): + """Save mid-training, load into new model+optimizer, verify training continues.""" + import os + import tempfile + + _set_quantization_test_seed(42) + hybrid_recipe = self._hybrid_checkpoint_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + x = torch.randn(4, 32, 256, dtype=torch.bfloat16, device="cuda") + target = torch.randn_like(x) + + for _ in range(3): + self._checkpoint_training_step(model, optimizer, x, target, hybrid_recipe) + + with tempfile.NamedTemporaryFile(delete=False, suffix=".pt") as checkpoint_file: + torch.save( + { + "model": model.state_dict(), + "optimizer": optimizer.state_dict(), + "cpu_rng_state": torch.get_rng_state(), + "cuda_rng_state": torch.cuda.get_rng_state(), + }, + checkpoint_file.name, + ) + checkpoint_path = checkpoint_file.name + + try: + uninterrupted = self._checkpoint_training_step( + model, + optimizer, + x, + target, + hybrid_recipe, + ) + + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + resumed_model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + resumed_optimizer = te.optimizers.FusedAdam( + resumed_model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + checkpoint = torch.load(checkpoint_path, weights_only=False) + resumed_model.load_state_dict(checkpoint["model"]) + resumed_optimizer.load_state_dict(checkpoint["optimizer"]) + torch.set_rng_state(checkpoint["cpu_rng_state"]) + torch.cuda.set_rng_state(checkpoint["cuda_rng_state"]) + + resumed = self._checkpoint_training_step( + resumed_model, + resumed_optimizer, + x, + target, + hybrid_recipe, + ) + + _assert_nested_state_exact( + resumed, + uninterrupted, + path="checkpoint continuation", + ) + finally: + os.unlink(checkpoint_path) + + +# --------------------------------------------------------------------------- +# 11. Activation recomputation (torch.utils.checkpoint / te.checkpoint) +# --------------------------------------------------------------------------- + + +def _reset_rng(seed: int = 1234): + """Reset deterministic RNG for reproducible forward/backward comparisons. + + Activation recompute relies on RNG equality between the first forward + and the recomputed forward. These tests use dropout-free modules, so + RNG advancement doesn't affect numerics, but we still reset between + runs so the reference (no-recompute) and checkpointed paths see + identical weight init, input, and grad_output seeds. + """ + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def _collect_outputs(out, inp, model): + """Gather forward output, input grad, and parameter grads into a flat list. + + Mirrors ``test_numerics.py::_test_e2e_*_recompute`` conventions so the + comparison against a non-recomputed baseline is a simple zip. + """ + results = [out.detach().clone()] + if inp.grad is not None: + results.append(inp.grad.detach().clone()) + for _, p in model.named_parameters(): + if p.requires_grad and p.grad is not None: + results.append(p.grad.detach().clone()) + return results + + +def _assert_outputs_bitwise_equal(ref, test, label): + """All stateless same-format hybrid recipes should be bitwise-identical + under activation recompute: same input bytes → same quantized bytes → + same GEMM result. Any drift means the recompute path silently diverged + (e.g. fell back to a different quantization path).""" + assert len(ref) == len(test), f"{label}: output count mismatch" + for i, (r, t) in enumerate(zip(ref, test)): + torch.testing.assert_close( + t, r, rtol=0, atol=0, msg=f"{label}: bitwise mismatch at output {i}" + ) + + +@requires_fp8 +class TestHybridActivationRecompute: + """Activation recomputation around TE modules with a hybrid CustomRecipe. + + Probes the interaction between ``HybridQuantizedTensor`` / + ``HybridQuantizedTensorStorage`` and the three activation-checkpoint + paths in use today: + + * ``te.checkpoint(fn, ..., use_reentrant=True)`` — reentrant path; wraps + ``torch.autograd.Function`` that re-runs the forward under + ``activation_recompute_forward(recompute_phase=True)``. This is the + Megatron-style path. + * ``te.checkpoint(fn, ..., use_reentrant=False)`` — non-reentrant path; + uses ``_checkpoint_hook`` (torch saved-tensors hooks) to discard + saved tensors on the first forward and recompute them on unpack. + * ``torch.utils.checkpoint.checkpoint(fn, ..., use_reentrant=False)`` + — vanilla PyTorch path without TE wrapper. Exercised because users + (and some Megatron configs) invoke it directly around TE modules. + + Failure modes it catches: + + * Silent BF16 fallback during recompute (would break bitwise parity + but pass loose tolerance — hence the bitwise assertion for + same-format stateless recipes). + * ``HybridQuantizedTensorStorage.prepare_for_saving`` / + ``restore_from_saved`` chain losing a sub-storage across the + save-for-backward boundary. + * ``HybridQuantizedTensor`` subclass being stripped by the autograd + engine (would manifest as ``AttributeError`` on the recomputed + tensor). + """ + + in_features = 128 + out_features = 128 + batch = 32 + + # ----- helpers --------------------------------------------------- + + def _same_format_fp8_recipe(self): + """Same-format FP8 current scaling both directions → bitwise-safe + baseline. Matches + :class:`TestHybridGemmBitwiseIdentical` construction so + recompute parity can be asserted bitwise-equal.""" + return _hybrid_custom_recipe( + row_factory=_fp8_row_factory, + col_factory=_fp8_col_factory, + grad_factory=_fp8_grad_factory, + ) + + def _same_format_mxfp8_recipe(self): + """Same-format MXFP8 both directions — stateless, per-block scales + computed from the tensor content; bitwise-stable under recompute.""" + return _hybrid_custom_recipe( + row_factory=_mxfp8_factory, + col_factory=_mxfp8_factory, + grad_factory=lambda: MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E5M2), + ) + + def _cross_format_fp8_mxfp8_recipe(self): + """Cross-format FP8 row + MXFP8 col — the canonical hybrid + scenario. Numerical parity is not bitwise because the wgrad GEMM + uses MXFP8 scaling modes on both operands (so grad_output must be + MXFP8 columnwise), pairing differently from the fprop path.""" + return _hybrid_custom_recipe( + row_factory=_fp8_row_factory, + col_factory=_mxfp8_factory, + grad_factory=_mxfp8_factory, + ) + + def _run_linear(self, recipe_obj, *, checkpoint_fn=None): + """Build a fresh Linear, run forward+backward, return collected + outputs. ``checkpoint_fn`` is an optional callable of the form + ``fn(model, inp) -> output`` that wraps the forward in an + activation-checkpoint implementation; ``None`` is the reference + (non-recompute) baseline. + """ + _reset_rng(seed=4242) + model = Linear(self.in_features, self.out_features, params_dtype=torch.bfloat16).cuda() + inp = torch.randn( + self.batch, + self.in_features, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + inp.retain_grad() + + with autocast(enabled=True, recipe=recipe_obj): + out = checkpoint_fn(model, inp) if checkpoint_fn is not None else model(inp) + out.float().sum().backward() + return _collect_outputs(out, inp, model) + + def _run_transformer_layer(self, recipe_obj, *, checkpoint_fn=None): + """Small TransformerLayer (no dropout, fuse_qkv) with optional + activation checkpointing around the whole block.""" + _reset_rng(seed=5151) + hidden = 128 + ffn = 128 + nheads = 4 + seq = 8 + bs = 4 + + model = TransformerLayer( + hidden, + ffn, + nheads, + hidden_dropout=0.0, + attention_dropout=0.0, + fuse_qkv_params=True, + params_dtype=torch.bfloat16, + ).cuda() + + inp = torch.randn(seq, bs, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True) + inp.retain_grad() + + with autocast(enabled=True, recipe=recipe_obj): + out = checkpoint_fn(model, inp) if checkpoint_fn is not None else model(inp) + out.float().sum().backward() + return _collect_outputs(out, inp, model) + + # ----- te.checkpoint, reentrant --------------------------------- + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_te_checkpoint_reentrant_linear_fp8_bitwise(self): + """te.checkpoint(use_reentrant=True) around te.Linear with + same-format FP8 hybrid → bitwise parity with non-recompute. + + This is the Megatron-style activation-recompute path. Bitwise + parity catches silent BF16 fallback (would pass loose tolerance). + """ + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=True) + + ref = self._run_linear(self._same_format_fp8_recipe(), checkpoint_fn=None) + test = self._run_linear(self._same_format_fp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(reentrant) FP8") + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_te_checkpoint_reentrant_linear_mxfp8_bitwise(self): + """Same as FP8 but MXFP8 hybrid — per-block scales must recompute + identically. Asserts that the MXFP8 path does not get disabled + during recompute.""" + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=True) + + ref = self._run_linear(self._same_format_mxfp8_recipe(), checkpoint_fn=None) + test = self._run_linear(self._same_format_mxfp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(reentrant) MXFP8") + + # ----- te.checkpoint, non-reentrant ----------------------------- + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_te_checkpoint_non_reentrant_linear_fp8_bitwise(self): + """te.checkpoint(use_reentrant=False) — the saved-tensors-hooks + path. Different recompute infra (``_checkpoint_hook``) than the + reentrant path; validates the hybrid activation survives the + pack/unpack transport.""" + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=False) + + ref = self._run_linear(self._same_format_fp8_recipe(), checkpoint_fn=None) + test = self._run_linear(self._same_format_fp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(non-reentrant) FP8") + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_te_checkpoint_non_reentrant_linear_mxfp8_bitwise(self): + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=False) + + ref = self._run_linear(self._same_format_mxfp8_recipe(), checkpoint_fn=None) + test = self._run_linear(self._same_format_mxfp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(non-reentrant) MXFP8") + + # ----- torch.utils.checkpoint (vanilla, non-reentrant) ---------- + # + # These tests document a *known* TE-level incompatibility between + # vanilla ``torch.utils.checkpoint.checkpoint(..., use_reentrant=False)`` + # and TE's weight-workspace cache (``_linear_forward_impl`` in + # ``module/linear.py``). The mechanism: + # + # * First forward: ``quantize_weight`` takes the cache-miss path, + # creating a fresh hybrid workspace and threading it into + # ``prepare_for_saving`` → ``ctx.save_for_backward``. + # * Recompute forward: the workspace is already populated on the + # module, so ``quantize_weight`` takes the cache-hit path and + # saves a different tensor-count. + # + # Vanilla ``torch.utils.checkpoint`` (``use_reentrant=False``) + # enforces a strict count match between original-forward and + # recompute-forward ``save_for_backward`` calls, and rejects the + # discrepancy with ``CheckpointError: A different number of tensors + # was saved``. The 2:1 count ratio (``8`` forward vs ``4`` recompute) + # is a hybrid signature — both sub-storages are saved on cache-miss + # and only the remaining one on cache-hit. + # + # ``te.checkpoint`` avoids this by threading ``is_first_microbatch`` + # / ``skip_fp8_weight_update`` correctly across the recompute phase, + # which is why the ``te.checkpoint`` tests above pass bitwise. + # + # Keeping the xfail'd tests here (tracked by #3158): + # 1. pins the boundary — users hitting this failure get a clear + # diagnosis and pointer to ``te.checkpoint``; + # 2. becomes a regression signal if the underlying cache-vs- + # checkpoint interaction is ever resolved (the xfail flips to + # an unexpected pass). + # + # Not hybrid-specific *in nature* (any quantized TE module with + # weight-workspace caching hits it under vanilla torch checkpoint), + # but hybrid amplifies and surfaces it via the 2x sub-storage count. + + _TORCH_CHECKPOINT_FP8_XFAIL = pytest.mark.xfail( + raises=( + NotImplementedError + if not is_non_tn_fp8_gemm_supported() + else torch.utils.checkpoint.CheckpointError + ), + strict=True, + reason=( + "On Hopper, same-format FP8 hybrid reaches the unsupported columnwise-only " + "per-tensor quantization guard before checkpointing. On architectures that " + "support non-TN FP8 GEMMs, vanilla torch.utils.checkpoint(use_reentrant=False) " + "is incompatible with TE's weight-workspace cache. Tracked by #3158." + ), + ) + + _TORCH_CHECKPOINT_CACHE_XFAIL = pytest.mark.xfail( + raises=torch.utils.checkpoint.CheckpointError, + strict=True, + reason=( + "Vanilla torch.utils.checkpoint(use_reentrant=False) is incompatible " + "with TE's weight-workspace cache. Tracked by #3158." + ), + ) + + @_TORCH_CHECKPOINT_FP8_XFAIL + def test_torch_checkpoint_non_reentrant_linear_fp8_bitwise(self): + """Vanilla ``torch.utils.checkpoint.checkpoint`` without TE wrapper + around a hybrid-quantized te.Linear. + + Users invoke ``torch.utils.checkpoint`` directly in many Megatron + branches and custom recomputation schemes. Currently fails due to + the weight-workspace cache interaction documented above; pins the + boundary so a future fix would flip this to an unexpected pass. + """ + + def fn(model, inp): + return torch.utils.checkpoint.checkpoint(model, inp, use_reentrant=False) + + ref = self._run_linear(self._same_format_fp8_recipe(), checkpoint_fn=None) + test = self._run_linear(self._same_format_fp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "torch.utils.checkpoint FP8") + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + @_TORCH_CHECKPOINT_CACHE_XFAIL + def test_torch_checkpoint_non_reentrant_linear_mxfp8_bitwise(self): + def fn(model, inp): + return torch.utils.checkpoint.checkpoint(model, inp, use_reentrant=False) + + ref = self._run_linear(self._same_format_mxfp8_recipe(), checkpoint_fn=None) + test = self._run_linear(self._same_format_mxfp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "torch.utils.checkpoint MXFP8") + + # ----- cross-format + recompute (functional, loose tolerance) --- + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_te_checkpoint_reentrant_linear_cross_format(self): + """Cross-format hybrid (FP8 row + MXFP8 col) under activation + recompute. Numerics are allowed to drift from non-recompute only + through paths recompute is allowed to affect; in practice they + should still match tightly because the recipe is stateless. Loose + tolerance catches only catastrophic silent fallbacks.""" + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=True) + + ref = self._run_linear(self._cross_format_fp8_mxfp8_recipe(), checkpoint_fn=None) + test = self._run_linear(self._cross_format_fp8_mxfp8_recipe(), checkpoint_fn=fn) + # Expected to match bitwise since both quantizers are stateless + # and the input bytes are identical between the two runs. Use a + # strict tolerance; if this ever drifts it's a real bug. + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(reentrant) FP8xMXFP8 cross-format") + + # ----- TransformerLayer ----------------------------------------- + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_te_checkpoint_reentrant_transformer_layer_fp8(self): + """te.checkpoint(reentrant) around a full TransformerLayer under + hybrid FP8. Exercises LayerNormLinear + DPA + LayerNormMLP in one + shot — the ``with_quantized_norm=False`` unfused path for hybrid + in ``layernorm_linear.py`` / ``layernorm_mlp.py`` must produce + the same result when recomputed. + + Asserted bitwise: the module uses ``hidden_dropout=0.0``, + ``attention_dropout=0.0``, and ``te.checkpoint`` restores RNG + state before recompute, so every kernel sees identical inputs and + there are no stochastic ops. Non-determinism at this level would + indicate a real regression (e.g. a kernel quietly taking a + non-deterministic code path) — not measurement noise.""" + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=True) + + ref = self._run_transformer_layer(self._same_format_fp8_recipe(), checkpoint_fn=None) + test = self._run_transformer_layer(self._same_format_fp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(reentrant) TransformerLayer FP8") + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_te_checkpoint_non_reentrant_transformer_layer_fp8(self): + """Same TransformerLayer setup but through the non-reentrant + saved-tensors-hooks recompute path. Same bitwise-equality + rationale as the reentrant variant above.""" + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=False) + + ref = self._run_transformer_layer(self._same_format_fp8_recipe(), checkpoint_fn=None) + test = self._run_transformer_layer(self._same_format_fp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal( + ref, test, "te.checkpoint(non-reentrant) TransformerLayer FP8" + ) + + # ----- quantized_model_init + recompute ------------------------- + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_te_checkpoint_reentrant_quantized_model_init_fp8_bitwise(self): + """Combine ``quantized_model_init`` (persistent + HybridQuantizedTensor weights) with activation recompute — + verifies the recompute path doesn't try to re-quantize an already- + quantized weight incorrectly, and the HybridQuantizer workspace + caching stays consistent across first-forward + recomputed-forward.""" + import transformer_engine.pytorch as te_pytorch + + hybrid_recipe = self._same_format_fp8_recipe() + + def _build_and_run(use_checkpoint): + _reset_rng(seed=7777) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear( + self.in_features, self.out_features, params_dtype=torch.bfloat16 + ).cuda() + inp = torch.randn( + self.batch, + self.in_features, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + inp.retain_grad() + with autocast(enabled=True, recipe=hybrid_recipe): + if use_checkpoint: + out = te_pytorch.checkpoint(model, inp, use_reentrant=True) + else: + out = model(inp) + out.float().sum().backward() + return _collect_outputs(out, inp, model) + + ref = _build_and_run(use_checkpoint=False) + test = _build_and_run(use_checkpoint=True) + _assert_outputs_bitwise_equal( + ref, test, "quantized_model_init + te.checkpoint(reentrant) FP8" + ) + + # ----- GroupedLinear + recompute -------------------------------- + + def _run_grouped_linear(self, recipe_obj, *, checkpoint_fn=None): + """Build a GroupedLinear, run forward+backward with optional + activation checkpointing around the module. Exercises the + ``_split_quantize_hybrid`` code path under recompute. + + GroupedLinear is the MoE token-dispatch kernel: a single batch + is split along dim-0 into ``num_gemms`` chunks and each chunk + goes through its own weight matrix. Under hybrid quantization, + ``_split_quantize_hybrid`` (``module/grouped_linear.py``) runs + ``tex.split_quantize`` twice (once per sub-quantizer direction) + and zips the results into a list of ``HybridQuantizedTensor`` + chunks — save-for-backward then receives a *list* of hybrid + tensors, not a single one, so the ``prepare_for_saving`` chain + has to handle an extended tensor-object list. + """ + _reset_rng(seed=9090) + num_gemms = 3 + hidden = 128 + ffn = 128 + bs = 24 + + model = GroupedLinear(num_gemms, hidden, ffn, params_dtype=torch.bfloat16).cuda() + inp = torch.randn(bs, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True) + inp.retain_grad() + base = bs // num_gemms + rem = bs % num_gemms + m_splits = [base + (1 if i < rem else 0) for i in range(num_gemms)] + + with autocast(enabled=True, recipe=recipe_obj): + if checkpoint_fn is not None: + out = checkpoint_fn(model, inp, m_splits) + else: + out = model(inp, m_splits) + out.float().sum().backward() + return _collect_outputs(out, inp, model) + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_te_checkpoint_reentrant_grouped_linear_fp8_bitwise(self): + """GroupedLinear + te.checkpoint(reentrant) under same-format FP8 + hybrid. Exercises the MoE ``_split_quantize_hybrid`` + list-of- + hybrid-tensors save-for-backward path under recompute.""" + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp, m_splits): + return te_pytorch.checkpoint(model, inp, m_splits, use_reentrant=True) + + ref = self._run_grouped_linear(self._same_format_fp8_recipe(), checkpoint_fn=None) + test = self._run_grouped_linear(self._same_format_fp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(reentrant) GroupedLinear FP8") + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_te_checkpoint_non_reentrant_grouped_linear_fp8_bitwise(self): + """Same GroupedLinear recompute setup but through the non- + reentrant saved-tensors-hooks path — verifies that the list of + hybrid activations survives the pack/unpack transport (one hook + invocation per split × per sub-storage buffer, not just one).""" + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp, m_splits): + return te_pytorch.checkpoint(model, inp, m_splits, use_reentrant=False) + + ref = self._run_grouped_linear(self._same_format_fp8_recipe(), checkpoint_fn=None) + test = self._run_grouped_linear(self._same_format_fp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(non-reentrant) GroupedLinear FP8") + + # ----- Selective attention recompute ---------------------------- + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + @_XFAIL_SELECTIVE_ATTENTION_RECOMPUTE + def test_selective_attention_recompute_transformer_layer_fp8_bitwise(self): + """``TransformerLayer(..., checkpoint_core_attention=True)`` — + the Megatron default memory-savings pattern. + + Unlike full-layer recompute (``te.checkpoint(layer, inp)``), + selective attention recompute is a TransformerLayer-internal + option: only the DPA (dot-product attention) block is wrapped + in a checkpoint, everything else runs normally. This is a + *different* code path in ``transformer.py`` from the + ``te.checkpoint(...)`` tests above — DPA internally invokes its + own checkpoint context around the attention kernel. + + For hybrid, the question is whether a hybrid activation produced + by LayerNormLinear (QKV projection) survives the DPA-internal + recompute boundary (which saves it for backward) and is + consumable by the backward GEMM unchanged. + + Bitwise because the model uses ``hidden_dropout=0.0``, + ``attention_dropout=0.0``, and the DPA checkpoint restores RNG + state — so reference and recomputed paths should be identical + to the last bit.""" + _reset_rng(seed=5151) + hidden = 128 + ffn = 128 + nheads = 4 + seq = 8 + bs = 4 + + def _run(checkpoint_core_attention): + _reset_rng(seed=5151) + model = TransformerLayer( + hidden, + ffn, + nheads, + hidden_dropout=0.0, + attention_dropout=0.0, + fuse_qkv_params=True, + params_dtype=torch.bfloat16, + ).cuda() + inp = torch.randn( + seq, bs, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) + inp.retain_grad() + with autocast(enabled=True, recipe=self._same_format_fp8_recipe()): + out = model(inp, checkpoint_core_attention=checkpoint_core_attention) + out.float().sum().backward() + return _collect_outputs(out, inp, model) + + ref = _run(checkpoint_core_attention=False) + test = _run(checkpoint_core_attention=True) + _assert_outputs_bitwise_equal(ref, test, "checkpoint_core_attention TransformerLayer FP8") + + # ----- Linear bitwise parametrized across all 4 stateless formats ----- + + @pytest.mark.parametrize( + "format_name,reentrant", + [ + pytest.param( + "fp8_current", + True, + id="fp8_current-reentrant", + marks=_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8, + ), + pytest.param( + "fp8_current", + False, + id="fp8_current-nonreentrant", + marks=_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8, + ), + pytest.param( + "mxfp8", + True, + id="mxfp8-reentrant", + marks=pytest.mark.skipif( + not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}" + ), + ), + pytest.param( + "mxfp8", + False, + id="mxfp8-nonreentrant", + marks=pytest.mark.skipif( + not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}" + ), + ), + pytest.param( + "block_fp8", + True, + id="block_fp8-reentrant", + marks=pytest.mark.skipif( + not fp8_block_scaling_available, + reason=f"BlockFP8: {reason_for_no_fp8_block_scaling}", + ), + ), + pytest.param( + "block_fp8", + False, + id="block_fp8-nonreentrant", + marks=pytest.mark.skipif( + not fp8_block_scaling_available, + reason=f"BlockFP8: {reason_for_no_fp8_block_scaling}", + ), + ), + pytest.param( + "nvfp4", + True, + id="nvfp4-reentrant", + marks=pytest.mark.skipif( + not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}" + ), + ), + pytest.param( + "nvfp4", + False, + id="nvfp4-nonreentrant", + marks=pytest.mark.skipif( + not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}" + ), + ), + ], + ) + def test_te_checkpoint_linear_all_stateless_formats_bitwise(self, format_name, reentrant): + """Bitwise parity of Linear + te.checkpoint across all four + stateless hybrid formats (FP8 current, MXFP8, BlockFP8, NVFP4), + both reentrant and non-reentrant. + + Each format has a distinct history of columnwise-only kernel + support — BlockFP8 required C++ null-check patches before + columnwise-only mode worked, NVFP4 has packed FP4 layout plus + optional RHT cache, MXFP8 has [128,4]/[4,128] scale padding. + The recompute path exercises columnwise-only sub-quantizers + (rowwise is freed after fprop and only recreated on backward), + so format-specific columnwise-only handling is on the critical + path. + + A regression in any of these would silently fall back to BF16 + during recompute; bitwise equality catches that immediately.""" + import transformer_engine.pytorch as te_pytorch + + row_factory, col_factory_for_grad, hw_skip, hw_reason = _QUANTIZER_CONFIGS[format_name] + # Most formats have a distinct E5M2 variant for grad; NVFP4 has + # only one format (col_factory_for_grad is None → reuse + # row_factory, which is what the existing hybrid NVFP4 tests do). + grad_factory = col_factory_for_grad if col_factory_for_grad is not None else row_factory + + hybrid_recipe = _hybrid_custom_recipe( + row_factory=row_factory, + col_factory=row_factory, + grad_factory=grad_factory, + ) + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=reentrant) + + ref = self._run_linear(hybrid_recipe, checkpoint_fn=None) + test = self._run_linear(hybrid_recipe, checkpoint_fn=fn) + label = ( + f"te.checkpoint({'reentrant' if reentrant else 'non-reentrant'}) Linear {format_name}" + ) + _assert_outputs_bitwise_equal(ref, test, label) + + # ----- save_for_backward round-trip (unit-level) ---------------- + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_prepare_restore_roundtrip_is_identity(self): + """Unit-level guarantee: the + ``prepare_for_saving`` / ``restore_from_saved`` chain used by + activation-recompute ``ctx.save_for_backward`` preserves both + sub-storages bitwise. + + This is the primitive the recompute path is built on; pinning it + here gives a focused failure signal independent of the module- + level recompute tests above.""" + torch.manual_seed(0) + inp = torch.randn(256, 256, dtype=torch.bfloat16, device="cuda") + hq = HybridQuantizer( + rowwise_quantizer=_fp8_row_factory(), + columnwise_quantizer=_fp8_col_factory(), + ) + hybrid = hq.quantize(inp) + expected = hybrid.dequantize() + + saved_tensors, saved_obj = hybrid.prepare_for_saving() + # Mimic the autograd ctx round-trip: all saved tensors pass + # through ``ctx.save_for_backward`` (a no-op for semantics). + leftover = saved_obj.restore_from_saved(list(saved_tensors)) + assert leftover == [], "restore_from_saved should consume every element" + torch.testing.assert_close(saved_obj.dequantize(), expected, rtol=0, atol=0) diff --git a/tests/pytorch/test_hybrid_quantization_fsdp2.py b/tests/pytorch/test_hybrid_quantization_fsdp2.py new file mode 100644 index 0000000000..cda3e7184a --- /dev/null +++ b/tests/pytorch/test_hybrid_quantization_fsdp2.py @@ -0,0 +1,1515 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for HybridQuantizedTensor behavior required by PyTorch FSDP2.""" + +import io + +import pytest +import torch + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex + +from hybrid_quantization_utils import ( + as_data_tensor_tuple as _as_data_tensor_tuple, + assert_hybrid_tensor_exact as _assert_hybrid_tensor_exact, + assert_storage_data_exact as _assert_storage_data_exact, + fp8_e4m3_factory as _fp8_row_factory, + fp8_e5m2_factory as _fp8_grad_factory, + hybrid_block_fp8_e4m3_qfactory as _hybrid_block_fp8_qfactory, + hybrid_custom_recipe as _hybrid_custom_recipe, + make_fp8_quantizer as _make_fp8_quantizer, + make_hybrid_quantizer_fp8_row_fp4_col as _make_hybrid_quantizer_fp8_row_fp4_col, + mxfp8_e4m3_factory as _mxfp8_factory, +) +from transformer_engine.common import recipe +from transformer_engine.pytorch import ( + Float8BlockQuantizer, + Float8CurrentScalingQuantizer, + Float8Quantizer, + Float8Tensor, + Float8TensorStorage, + HybridQuantizedTensor, + HybridQuantizer, + IdentityQuantizer, + Linear, + MXFP8Quantizer, + NVFP4Quantizer, + QuantizedTensor, + quantized_model_init, +) +from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported + +_fp8_col_factory = _fp8_row_factory + + +fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) +nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) +mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) +fp8_block_scaling_available, reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available( + return_reason=True +) + +_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 = pytest.mark.xfail( + condition=not is_non_tn_fp8_gemm_supported(), + raises=NotImplementedError, + strict=True, + reason=( + "Hopper does not yet support columnwise-only per-tensor FP8 quantization; " + "tracked by NVIDIA/TransformerEngine#3158" + ), +) + +requires_fp8 = pytest.mark.skipif( + not fp8_available, + reason=f"FP8: {reason_for_no_fp8}", +) + +requires_fp8_and_nvfp4 = pytest.mark.skipif( + not (fp8_available and nvfp4_available), + reason=f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", +) + +requires_mxfp8 = pytest.mark.skipif( + not mxfp8_available, + reason=f"MXFP8: {reason_for_no_mxfp8}", +) + + +# --------------------------------------------------------------------------- +# 1. __torch_dispatch__ operations that FSDP2 relies on +# --------------------------------------------------------------------------- + +aten = torch.ops.aten + + +def _make_hybrid_param_for_dispatch( + row_factory, col_factory, grad_factory=None, in_features=256, out_features=256 +): + """Create a HybridQuantizedTensor weight via quantized_model_init for dispatch tests.""" + hybrid_recipe = _hybrid_custom_recipe(row_factory, col_factory, grad_factory) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + return model.weight + + +_dispatch_configs = [ + pytest.param( + "fp8_fp8", + id="same-format-fp8", + marks=_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8, + ), +] +if mxfp8_available: + _dispatch_configs.append(pytest.param("mxfp8_mxfp8", id="same-format-mxfp8")) + + +def _get_dispatch_hybrid_param(config_name): + """Return a HybridQuantizedTensor weight for the given config.""" + if config_name == "fp8_fp8": + return _make_hybrid_param_for_dispatch( + _fp8_row_factory, + _fp8_col_factory, + _fp8_grad_factory, + ) + elif config_name == "mxfp8_mxfp8": + return _make_hybrid_param_for_dispatch( + _mxfp8_factory, + _mxfp8_factory, + grad_factory=lambda: MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E5M2), + ) + else: + raise ValueError(f"Unknown config: {config_name}") + + +@requires_fp8 +class TestFloat8TransposeOnlySplit: + """Regression coverage for columnwise-only Float8 split metadata. + + A columnwise-only per-tensor Float8 sub-storage may have ``_data=None`` and + store its bytes in ``_transpose`` with physical shape ``[K, M]``. Splitting + that tensor must still produce pieces whose wrapper shape is the logical + row-major shape ``[M_i, K]``; otherwise HybridQuantizedTensor uses the + transposed shape when rowwise storage is absent. + """ + + @staticmethod + def _make_transpose_only_float8_tensor(shape=(12, 16)): + m, k = shape + data_transpose = torch.empty((k, m), dtype=torch.uint8, device="cuda") + return Float8Tensor( + shape=shape, + dtype=torch.bfloat16, + data=None, + data_transpose=data_transpose, + fp8_scale_inv=torch.ones(1, dtype=torch.float32, device="cuda"), + fp8_dtype=tex.DType.kFloat8E4M3, + requires_grad=False, + device="cuda", + ) + + @pytest.mark.parametrize( + "split_size,dim,expected_shapes,expected_transpose_shapes", + [ + (5, 0, [(5, 16), (5, 16), (2, 16)], [(16, 5), (16, 5), (16, 2)]), + (6, 1, [(12, 6), (12, 6), (12, 4)], [(6, 12), (6, 12), (4, 12)]), + ], + ) + def test_float8_split_uses_logical_shape_for_transpose_only_storage( + self, split_size, dim, expected_shapes, expected_transpose_shapes + ): + tensor = self._make_transpose_only_float8_tensor() + + pieces = torch.split(tensor, split_size, dim=dim) + + assert [tuple(piece.shape) for piece in pieces] == expected_shapes + assert [tuple(piece._transpose.shape) for piece in pieces] == expected_transpose_shapes + assert all(piece._data is None for piece in pieces) + assert all(piece._transpose_invalid is False for piece in pieces) + + def test_float8_view_preserves_transpose_only_storage(self): + tensor = self._make_transpose_only_float8_tensor() + + viewed = tensor.view(12, 16) + + assert tuple(viewed.shape) == (12, 16) + assert viewed._data is None + assert viewed._transpose is not None + assert tuple(viewed._transpose.shape) == (16, 12) + assert viewed._transpose_invalid is False + + def test_float8_aten_view_preserves_transpose_only_storage(self): + tensor = self._make_transpose_only_float8_tensor() + + viewed = torch.ops.aten.view.default(tensor, [12, 16]) + + assert tuple(viewed.shape) == (12, 16) + assert viewed._data is None + assert viewed._transpose is not None + assert tuple(viewed._transpose.shape) == (16, 12) + assert viewed._transpose_invalid is False + + def test_float8_storage_view_preserves_transpose_only_storage(self): + data_transpose = torch.empty((16, 12), dtype=torch.uint8, device="cuda") + storage = Float8TensorStorage( + data=None, + data_transpose=data_transpose, + fp8_scale_inv=torch.ones(1, dtype=torch.float32, device="cuda"), + fp8_dtype=tex.DType.kFloat8E4M3, + fake_dtype=torch.bfloat16, + ) + + viewed = storage.view(torch.Size((12, 16))) + + assert viewed._data is None + assert viewed._transpose is not None + assert tuple(viewed._transpose.shape) == (16, 12) + assert viewed._transpose_invalid is False + + def test_float8_shape_changing_view_raises_for_transpose_only_storage(self): + tensor = self._make_transpose_only_float8_tensor() + + with pytest.raises(NotImplementedError, match="columnwise-only data"): + tensor.view(6, 32) + + def test_hybrid_split_uses_columnwise_logical_shape_when_rowwise_is_absent(self): + columnwise = self._make_transpose_only_float8_tensor() + quantizer = HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + quantizer.set_usage(rowwise=False, columnwise=True) + hybrid = HybridQuantizedTensor( + shape=columnwise.shape, + dtype=columnwise.dtype, + rowwise_storage=None, + columnwise_storage=columnwise, + quantizer=quantizer, + device="cuda", + ) + + pieces = torch.split(hybrid, 5, dim=0) + + assert [tuple(piece.shape) for piece in pieces] == [ + (5, 16), + (5, 16), + (2, 16), + ] + assert all(piece.rowwise_sub_storage is None for piece in pieces) + assert [tuple(piece.columnwise_sub_storage.shape) for piece in pieces] == [ + (5, 16), + (5, 16), + (2, 16), + ] + assert [tuple(piece.columnwise_sub_storage._transpose.shape) for piece in pieces] == [ + (16, 5), + (16, 5), + (16, 2), + ] + + @staticmethod + def _make_valid_transpose_only_float8_tensor(shape=(12, 16)): + """Build transpose-only storage with known, numerically valid FP8 bytes.""" + source = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + rowwise_quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + rowwise=True, + columnwise=False, + ) + rowwise = rowwise_quantizer(source) + columnwise_quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + rowwise=False, + columnwise=True, + ) + tensor = Float8Tensor( + shape=shape, + dtype=torch.bfloat16, + data=None, + data_transpose=rowwise._data.movedim(-1, 0).contiguous(), + fp8_scale_inv=rowwise._scale_inv.detach().clone(), + fp8_dtype=tex.DType.kFloat8E4M3, + quantizer=columnwise_quantizer, + requires_grad=False, + device="cuda", + ) + return tensor, rowwise.dequantize() + + @pytest.mark.parametrize("weights_only", (False, True)) + def test_serialization_preserves_transpose_only_payload(self, weights_only): + tensor, expected = self._make_valid_transpose_only_float8_tensor() + buffer = io.BytesIO() + + torch.save(tensor, buffer) + buffer.seek(0) + loaded = torch.load(buffer, weights_only=weights_only) + + assert loaded._data is None + assert loaded._transpose_invalid is False + assert torch.equal(loaded._transpose, tensor._transpose) + torch.testing.assert_close(loaded.dequantize(), expected, rtol=0.0, atol=0.0) + + def test_dequantize_from_transpose_only_payload(self): + tensor, expected = self._make_valid_transpose_only_float8_tensor() + + torch.testing.assert_close(tensor.dequantize(), expected, rtol=0.0, atol=0.0) + + def test_slice_and_select_preserve_transpose_only_payload(self): + tensor, expected = self._make_valid_transpose_only_float8_tensor() + + sliced = tensor[2:8] + selected = torch.select(tensor, 1, 3) + + assert isinstance(sliced, Float8Tensor) + assert sliced._data is None + assert sliced._transpose.shape == torch.Size((16, 6)) + torch.testing.assert_close(sliced.dequantize(), expected[2:8], rtol=0.0, atol=0.0) + assert isinstance(selected, Float8Tensor) + assert selected._data is None + assert selected._transpose.shape == torch.Size((12,)) + torch.testing.assert_close(selected.dequantize(), expected[:, 3], rtol=0.0, atol=0.0) + + def test_as_strided_row_shard_preserves_transpose_only_payload(self): + tensor, expected = self._make_valid_transpose_only_float8_tensor() + + shard = torch.as_strided(tensor, (5, 16), (16, 1), 16) + + assert isinstance(shard, Float8Tensor) + assert shard._data is None + assert shard._transpose.shape == torch.Size((16, 5)) + torch.testing.assert_close(shard.dequantize(), expected[1:6], rtol=0.0, atol=0.0) + + def test_as_strided_falls_back_for_nonrepresentable_layout(self): + tensor, expected = self._make_valid_transpose_only_float8_tensor() + + output = torch.as_strided(tensor, (16, 12), (1, 16), 0) + + assert type(output) is torch.Tensor + reference = torch.as_strided(expected, (16, 12), (1, 16), 0) + torch.testing.assert_close(output, reference, rtol=0.0, atol=0.0) + + def test_new_zeros_preserves_transpose_only_layout(self): + tensor, _ = self._make_valid_transpose_only_float8_tensor() + + output = tensor.new_zeros((5, 16)) + + assert isinstance(output, Float8Tensor) + assert output.shape == torch.Size((5, 16)) + assert output._data is None + assert output._transpose.shape == torch.Size((16, 5)) + assert output._transpose.dtype == torch.uint8 + assert torch.count_nonzero(output._transpose).item() == 0 + torch.testing.assert_close( + output.dequantize(), + torch.zeros((5, 16), dtype=output.dtype, device=output.device), + rtol=0.0, + atol=0.0, + ) + + +class TestHybridNewZeros: + """Public new_zeros semantics independent of FSDP-specific allocation.""" + + @staticmethod + def _make_identity_hybrid(): + quantizer = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + source = quantizer(torch.ones((4, 8), dtype=torch.bfloat16, device="cuda")) + return quantizer, source + + @pytest.mark.parametrize( + "keep_rowwise,keep_columnwise", + [(True, True), (True, False), (False, True)], + ) + def test_initializes_every_present_direction_and_preserves_kwargs( + self, + keep_rowwise, + keep_columnwise, + ): + quantizer, source = self._make_identity_hybrid() + if not keep_rowwise: + source.update_usage(rowwise_usage=False) + if not keep_columnwise: + source.update_usage(columnwise_usage=False) + + # Direction selection must come from source storage, not mutable parent + # usage flags, and new_zeros must not change those flags. + quantizer.set_usage(rowwise=False, columnwise=False) + result = torch.ops.aten.new_zeros.default( + source, + [3, 5], + dtype=torch.float32, + device=source.device, + ) + + assert isinstance(result, HybridQuantizedTensor) + assert result.shape == torch.Size((3, 5)) + assert result.dtype == torch.float32 + assert result.device == source.device + assert (result.rowwise_sub_storage is not None) is keep_rowwise + assert (result.columnwise_sub_storage is not None) is keep_columnwise + assert quantizer.get_usages() == {"rowwise": False, "columnwise": False} + assert result._quantizer is not quantizer + assert result._quantizer.get_usages() == { + "rowwise": keep_rowwise, + "columnwise": keep_columnwise, + } + + for sub_storage in ( + result.rowwise_sub_storage, + result.columnwise_sub_storage, + ): + if sub_storage is not None: + torch.testing.assert_close( + sub_storage.dequantize(), + torch.zeros((3, 5), dtype=torch.float32, device=source.device), + rtol=0.0, + atol=0.0, + ) + + # A plain-source copy routes through the result's parent quantizer. It + # must update every allocated direction even though the source parent + # had both of its mutable usage flags disabled before new_zeros. + plain_source = torch.full((3, 5), 7.0, dtype=torch.float32, device=source.device) + result.copy_(plain_source) + for sub_storage in ( + result.rowwise_sub_storage, + result.columnwise_sub_storage, + ): + if sub_storage is not None: + torch.testing.assert_close( + sub_storage.dequantize(), plain_source, rtol=0.0, atol=0.0 + ) + assert quantizer.get_usages() == {"rowwise": False, "columnwise": False} + + def test_identity_substorages_allow_integer_dtype(self): + _, source = self._make_identity_hybrid() + result = source.new_zeros((2, 3), dtype=torch.int32) + + assert result.dtype == torch.int32 + for sub_storage in ( + result.rowwise_sub_storage, + result.columnwise_sub_storage, + ): + assert sub_storage.dequantize().dtype == torch.int32 + assert torch.count_nonzero(sub_storage.dequantize()).item() == 0 + + @requires_fp8 + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + @pytest.mark.parametrize("unsupported_dtype", (torch.int32, torch.float64, torch.bool)) + def test_non_identity_substorage_rejects_unsupported_dtype(self, unsupported_dtype): + quantizer = HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + source = quantizer(torch.ones((4, 8), dtype=torch.bfloat16, device="cuda")) + + with pytest.raises(TypeError, match="new_zeros only supports"): + source.new_zeros((2, 3), dtype=unsupported_dtype) + + def test_does_not_invoke_live_quantizers_or_consume_rng(self, monkeypatch): + quantizer, source = self._make_identity_hybrid() + cpu_rng_before = torch.get_rng_state().clone() + cuda_rng_before = torch.cuda.get_rng_state(source.device).clone() + + def fail_live_make_empty(*args, **kwargs): + raise AssertionError("new_zeros invoked a live quantizer") + + monkeypatch.setattr(quantizer, "make_empty", fail_live_make_empty) + monkeypatch.setattr( + source.rowwise_sub_storage._quantizer, + "make_empty", + fail_live_make_empty, + ) + monkeypatch.setattr( + source.columnwise_sub_storage._quantizer, + "make_empty", + fail_live_make_empty, + ) + + result = source.new_zeros((2, 6)) + + assert isinstance(result, HybridQuantizedTensor) + assert torch.equal(torch.get_rng_state(), cpu_rng_before) + assert torch.equal(torch.cuda.get_rng_state(source.device), cuda_rng_before) + + def test_rejects_empty_hybrid(self): + _, source = self._make_identity_hybrid() + source.update_usage(rowwise_usage=False, columnwise_usage=False) + + with pytest.raises(RuntimeError, match="at least one present sub-storage"): + source.new_zeros((2, 6)) + + @requires_fp8_and_nvfp4 + def test_initializes_nvfp4_data_scale_and_amax_buffers(self): + quantizer = _make_hybrid_quantizer_fp8_row_fp4_col() + source = quantizer(torch.randn((128, 256), dtype=torch.bfloat16, device="cuda")) + + result = source.new_zeros(source.shape) + + assert type(result.rowwise_sub_storage) is type(source.rowwise_sub_storage) + assert type(result.columnwise_sub_storage) is type(source.columnwise_sub_storage) + torch.testing.assert_close( + result.dequantize(), + torch.zeros_like(result.dequantize()), + rtol=0.0, + atol=0.0, + ) + for sub_storage in ( + result.rowwise_sub_storage, + result.columnwise_sub_storage, + ): + buffers, storage = sub_storage.prepare_for_saving() + try: + assert all( + buffer is None or torch.count_nonzero(buffer).item() == 0 for buffer in buffers + ) + finally: + assert storage.restore_from_saved(buffers) == [] + + @pytest.mark.skipif( + not fp8_block_scaling_available, + reason=f"Float8Blockwise: {reason_for_no_fp8_block_scaling}", + ) + def test_initializes_float8_block_data_and_scale_buffers(self): + quantizer = HybridQuantizer( + rowwise_quantizer=Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=2, + ), + columnwise_quantizer=Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=2, + ), + ) + source = quantizer(torch.randn((128, 256), dtype=torch.bfloat16, device="cuda")) + + result = source.new_zeros(source.shape) + + assert type(result.rowwise_sub_storage) is type(source.rowwise_sub_storage) + assert type(result.columnwise_sub_storage) is type(source.columnwise_sub_storage) + torch.testing.assert_close( + result.dequantize(), + torch.zeros_like(result.dequantize()), + rtol=0.0, + atol=0.0, + ) + for sub_storage in ( + result.rowwise_sub_storage, + result.columnwise_sub_storage, + ): + buffers, storage = sub_storage.prepare_for_saving() + try: + assert all( + buffer is None or torch.count_nonzero(buffer).item() == 0 for buffer in buffers + ) + finally: + assert storage.restore_from_saved(buffers) == [] + + +@requires_fp8 +class TestHybridTorchDispatchFSDP2Ops: + """Test aten ops that FSDP2 relies on to preserve the HybridQuantizedTensor type. + + Each op is called directly via torch.ops.aten and the result is verified to + still be HybridQuantizedTensor with valid sub-storages. + """ + + @pytest.fixture(params=_dispatch_configs) + def hybrid_param(self, request): + torch.manual_seed(42) + return _get_dispatch_hybrid_param(request.param) + + def test_split_preserves_hybrid_type(self, hybrid_param): + """torch.split must return a list of HybridQuantizedTensor pieces.""" + dim0 = hybrid_param.shape[0] + chunk_size = dim0 // 2 + pieces = torch.split(hybrid_param, chunk_size, dim=0) + expected_rowwise = torch.split(hybrid_param.rowwise_sub_storage, chunk_size, dim=0) + expected_columnwise = torch.split(hybrid_param.columnwise_sub_storage, chunk_size, dim=0) + assert len(pieces) == len(expected_rowwise) == len(expected_columnwise) + + assert len(pieces) >= 2 + for piece in pieces: + assert isinstance( + piece, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(piece).__name__}" + assert piece.rowwise_sub_storage is not None + assert piece.columnwise_sub_storage is not None + for index, (piece, expected_row, expected_column) in enumerate( + zip(pieces, expected_rowwise, expected_columnwise) + ): + _assert_storage_data_exact( + piece.rowwise_sub_storage, + expected_row, + context=f"split {index} rowwise", + ) + _assert_storage_data_exact( + piece.columnwise_sub_storage, + expected_column, + context=f"split {index} columnwise", + ) + + total_rows = sum(p.shape[0] for p in pieces) + assert total_rows == dim0 + + orig_deq = hybrid_param.dequantize() + reassembled = torch.cat([p.dequantize() for p in pieces], dim=0) + torch.testing.assert_close(orig_deq, reassembled, rtol=0.0, atol=0.0) + + def test_split_sub_storage_types_preserved(self, hybrid_param): + """After split, sub-storage types must match the original.""" + orig_row_type = type(hybrid_param.rowwise_sub_storage) + orig_col_type = type(hybrid_param.columnwise_sub_storage) + + chunk_size = hybrid_param.shape[0] // 2 + pieces = torch.split(hybrid_param, chunk_size, dim=0) + for piece in pieces: + assert type(piece.rowwise_sub_storage) is orig_row_type + assert type(piece.columnwise_sub_storage) is orig_col_type + + @requires_mxfp8 + def test_split_rejects_mxfp8_high_precision_fallback(self): + """Fail before wrapping unquantizable MXFP8 shards for Hybrid FSDP2.""" + quantizer = HybridQuantizer( + rowwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + columnwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + ) + tensor = quantizer(torch.randn(64, 64, dtype=torch.bfloat16, device="cuda")) + + with pytest.raises(NotImplementedError, match="local shape.*divisible by 32"): + torch.split(tensor, 16, dim=0) + + def test_view_preserves_hybrid_type(self, hybrid_param): + """view must return a HybridQuantizedTensor (used by FSDP2 reset_sharded_param).""" + shape_2d = hybrid_param.shape + result = aten.view.default(hybrid_param, list(shape_2d)) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + assert result.rowwise_sub_storage is not None + assert result.columnwise_sub_storage is not None + + def test_view_same_shape_preserves_hybrid(self, hybrid_param): + """view with same shape must return HybridQuantizedTensor.""" + shape_2d = list(hybrid_param.shape) + result = aten.view.default(hybrid_param, shape_2d) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + + def test_as_strided_noop_preserves_hybrid(self, hybrid_param): + """as_strided with matching shape/strides is a no-op that preserves type.""" + shape = tuple(hybrid_param.size()) + strides = (shape[-1], 1) + result = aten.as_strided.default(hybrid_param, list(shape), list(strides)) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + assert result.rowwise_sub_storage is not None + assert result.columnwise_sub_storage is not None + + def test_slice_noop_preserves_hybrid(self, hybrid_param): + """slice with full range is a no-op that preserves type.""" + result = aten.slice.Tensor(hybrid_param, 0, 0, hybrid_param.size(0)) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + assert result.rowwise_sub_storage is not None + + def test_copy_between_hybrid_tensors(self, hybrid_param): + """copy_ between compatible HybridQuantizedTensors copies quantized data directly.""" + src_deq = hybrid_param.dequantize().clone() + dst = hybrid_param._quantizer.make_empty( + shape=hybrid_param.shape, + dtype=hybrid_param.dtype, + device=hybrid_param.device, + ) + assert isinstance(dst, HybridQuantizedTensor) + + aten.copy_.default(dst, hybrid_param) + dst_deq = dst.dequantize() + torch.testing.assert_close(src_deq, dst_deq, rtol=0.0, atol=0.0) + _assert_hybrid_tensor_exact(dst, hybrid_param, context="hybrid copy_") + + def test_copy_between_mismatched_usages_raises_atomically(self, hybrid_param): + quantizer = hybrid_param._quantizer.copy() + src = quantizer.quantize(torch.ones_like(hybrid_param.dequantize())) + dst = quantizer.quantize(torch.zeros_like(hybrid_param.dequantize())) + src.update_usage(columnwise_usage=False) + dst_before = tuple( + None if tensor is None else tensor.clone() for tensor in _as_data_tensor_tuple(dst) + ) + + with pytest.raises( + NotImplementedError, + match="requires matching rowwise/columnwise usages", + ): + aten.copy_.default(dst, src) + + dst_after = _as_data_tensor_tuple(dst) + assert len(dst_after) == len(dst_before) + for before, after in zip(dst_before, dst_after): + if before is None: + assert after is None + else: + torch.testing.assert_close(after, before, rtol=0.0, atol=0.0) + + def test_copy_from_bf16_to_hybrid(self, hybrid_param): + """copy_ from BF16 into HybridQuantizedTensor triggers quantize_.""" + param = hybrid_param.detach() + bf16_data = torch.randn_like(param.dequantize()) + expected = param._quantizer.quantize(bf16_data) + aten.copy_.default(param, bf16_data) + assert isinstance(param, HybridQuantizedTensor) + _assert_hybrid_tensor_exact(param, expected, context="BF16 copy_") + torch.testing.assert_close(param.dequantize(), expected.dequantize(), rtol=0.0, atol=0.0) + + def test_new_zeros_returns_hybrid(self, hybrid_param): + """new_zeros returns initialized storage that remains FSDP-copyable.""" + new_shape = list(hybrid_param.shape) + result = aten.new_zeros.default(hybrid_param, new_shape) + + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + assert result.shape == hybrid_param.shape + assert result.rowwise_sub_storage is not None + assert result.columnwise_sub_storage is not None + assert type(result.rowwise_sub_storage) is type(hybrid_param.rowwise_sub_storage) + assert type(result.columnwise_sub_storage) is type(hybrid_param.columnwise_sub_storage) + torch.testing.assert_close( + result.dequantize(), + torch.zeros_like(result.dequantize()), + rtol=0.0, + atol=0.0, + ) + + # FSDP2 overwrites the initialized destination via copy_ after gather. + aten.copy_.default(result, hybrid_param) + torch.testing.assert_close( + result.dequantize(), hybrid_param.dequantize(), rtol=0.0, atol=0.0 + ) + _assert_hybrid_tensor_exact(result, hybrid_param, context="new_zeros then copy_") + + def test_empty_like_returns_hybrid(self, hybrid_param): + """empty_like must return a HybridQuantizedTensor.""" + result = aten.empty_like.default(hybrid_param) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + assert result.shape == hybrid_param.shape + assert result.rowwise_sub_storage is not None + + def test_clone_returns_hybrid(self, hybrid_param): + """clone must return an independent HybridQuantizedTensor with same data.""" + result = aten.clone.default(hybrid_param) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + assert result is not hybrid_param + torch.testing.assert_close( + result.dequantize(), hybrid_param.dequantize(), rtol=0.0, atol=0.0 + ) + _assert_hybrid_tensor_exact(result, hybrid_param, context="clone") + + +# --------------------------------------------------------------------------- +# 2. fsdp_pre_all_gather protocol +# --------------------------------------------------------------------------- + + +def _make_fsdp_protocol_param(config_name): + """Create a HybridQuantizedTensor weight for FSDP protocol tests.""" + if config_name == "fp8_fp8": + r = _hybrid_custom_recipe(_fp8_row_factory, _fp8_col_factory, _fp8_grad_factory) + elif config_name == "mxfp8_fp8": + r = _hybrid_custom_recipe(_mxfp8_factory, _fp8_col_factory, _fp8_grad_factory) + elif config_name == "block_fp8": + r = recipe.CustomRecipe(qfactory=_hybrid_block_fp8_qfactory) + else: + raise ValueError(f"Unknown config: {config_name}") + with quantized_model_init(enabled=True, recipe=r): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + return model.weight + + +_fsdp_protocol_configs = [ + pytest.param( + "fp8_fp8", + id="same-format", + marks=_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8, + ) +] +if mxfp8_available: + _fsdp_protocol_configs.append(pytest.param("mxfp8_fp8", id="mixed-mxfp8-fp8")) +if fp8_block_scaling_available: + _fsdp_protocol_configs.append(pytest.param("block_fp8", id="same-format-block-fp8")) + + +@requires_fp8 +class TestHybridFsdpPreAllGatherProtocol: + """Test the fsdp_pre_all_gather method on HybridQuantizedTensor. + + These tests call the method directly (no actual all-gather communication) + to verify the protocol contract: returns (sharded_tensors, metadata) where + sharded_tensors is a tuple of plain torch.Tensor. + """ + + @pytest.fixture(params=_fsdp_protocol_configs) + def hybrid_param(self, request): + torch.manual_seed(42) + return _make_fsdp_protocol_param(request.param) + + def test_pre_all_gather_returns_tuple_pair(self, hybrid_param): + """fsdp_pre_all_gather returns (sharded_tensors, metadata).""" + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + assert isinstance( + sharded_tensors, tuple + ), f"sharded_tensors should be tuple, got {type(sharded_tensors).__name__}" + assert len(sharded_tensors) > 0, "sharded_tensors should not be empty" + assert isinstance( + metadata, tuple + ), f"metadata should be tuple, got {type(metadata).__name__}" + + def test_pre_all_gather_buffers_are_plain_tensors(self, hybrid_param): + """Every element in sharded_tensors must be a plain torch.Tensor.""" + sharded_tensors, _ = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + for i, t in enumerate(sharded_tensors): + assert isinstance( + t, torch.Tensor + ), f"sharded_tensors[{i}] should be torch.Tensor, got {type(t).__name__}" + assert not isinstance( + t, QuantizedTensor + ), f"sharded_tensors[{i}] should NOT be QuantizedTensor subclass" + + def test_pre_all_gather_buffer_count_consistent(self, hybrid_param): + """Buffer count must be the same across repeated calls (FSDP2 buffer reuse).""" + sharded_1, _ = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + sharded_2, _ = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + assert len(sharded_1) == len( + sharded_2 + ), f"Buffer count changed: {len(sharded_1)} vs {len(sharded_2)}" + + def test_pre_all_gather_metadata_sufficient_for_reconstruction(self, hybrid_param): + """Metadata must contain enough info to reconstruct the tensor.""" + _, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + assert metadata is not None + assert len(metadata) > 0, "metadata should not be empty" + + +# --------------------------------------------------------------------------- +# 3. fsdp_post_all_gather protocol +# --------------------------------------------------------------------------- + + +@requires_fp8 +class TestHybridFsdpPostAllGatherProtocol: + """Test the fsdp_post_all_gather method on HybridQuantizedTensor. + + Simulates the post-all-gather phase by passing the sharded_tensors + from pre_all_gather directly (mimicking a single-rank all-gather). + """ + + @pytest.fixture(params=_fsdp_protocol_configs) + def hybrid_param(self, request): + torch.manual_seed(42) + return _make_fsdp_protocol_param(request.param) + + def test_post_all_gather_first_call_returns_hybrid_tensor(self, hybrid_param): + """With out=None, post_all_gather returns (HybridQuantizedTensor, outputs).""" + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + result, ag_outputs = hybrid_param.fsdp_post_all_gather( + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, + ) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + assert result.shape == hybrid_param.shape + assert result.rowwise_sub_storage is not None + assert result.columnwise_sub_storage is not None + + def test_post_all_gather_buffer_reuse(self, hybrid_param): + """On second call with out=previous, the same object is returned (buffer reuse).""" + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + first_result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, + ) + + second_result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, + metadata, + hybrid_param.dtype, + out=first_result, + ) + assert ( + second_result is first_result + ), "Buffer reuse: post_all_gather(out=prev) should return the same object" + _assert_hybrid_tensor_exact(second_result, hybrid_param, context="post-all-gather reuse") + + def test_post_all_gather_dequantize_matches_original(self, hybrid_param): + """Reconstructed tensor should dequantize close to the original.""" + orig_deq = hybrid_param.dequantize() + + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, + ) + result_deq = result.dequantize() + _assert_hybrid_tensor_exact(result, hybrid_param, context="post-all-gather") + torch.testing.assert_close(orig_deq, result_deq, rtol=0.0, atol=0.0) + + def test_post_all_gather_sub_storage_types_correct(self, hybrid_param): + """Reconstructed tensor's sub-storages match the original types.""" + orig_row_type = type(hybrid_param.rowwise_sub_storage) + orig_col_type = type(hybrid_param.columnwise_sub_storage) + + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, + ) + assert type(result.rowwise_sub_storage) is orig_row_type + _assert_hybrid_tensor_exact(result, hybrid_param, context="post-all-gather storage types") + assert type(result.columnwise_sub_storage) is orig_col_type + + +# --------------------------------------------------------------------------- +# 4. Pre/post all-gather roundtrip +# --------------------------------------------------------------------------- + + +@requires_fp8 +class TestHybridFsdpRoundtrip: + """End-to-end single-process roundtrip (pre -> post) without communication.""" + + @pytest.fixture(params=_fsdp_protocol_configs) + def hybrid_param(self, request): + torch.manual_seed(42) + return _make_fsdp_protocol_param(request.param) + + def test_pre_post_roundtrip_preserves_data(self, hybrid_param): + """pre_all_gather -> post_all_gather(out=None) -> dequantize matches original.""" + orig_deq = hybrid_param.dequantize() + + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, + ) + _assert_hybrid_tensor_exact(result, hybrid_param, context="FSDP roundtrip") + torch.testing.assert_close(orig_deq, result.dequantize(), rtol=0.0, atol=0.0) + + def test_pre_post_roundtrip_buffer_reuse_preserves_data(self, hybrid_param): + """Second roundtrip with out=previous preserves data (iteration 2+ simulation).""" + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + first_result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, + ) + + sharded_tensors_2, metadata_2 = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + second_result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors_2, + metadata_2, + hybrid_param.dtype, + out=first_result, + ) + assert second_result is first_result + torch.testing.assert_close( + hybrid_param.dequantize(), second_result.dequantize(), rtol=0.0, atol=0.0 + ) + _assert_hybrid_tensor_exact(second_result, hybrid_param, context="FSDP roundtrip reuse") + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_scale_refresh_across_iterations(self): + """After a sharded optimizer-style requantize, iter-2+ gathers see the new scale. + + Per-tensor FP8 does NOT include ``_scale_inv`` in ``fsdp_buffer_fields`` + (only ``_data`` is gathered; the scalar scale travels via iter-1 + metadata). This relies on the invariant that the sharded and gathered + ``Float8Tensor`` s share the same ``_scale_inv`` tensor object, and + that ``Float8CurrentScalingQuantizer.update_quantized`` writes the new + scale in place rather than replacing the tensor reference. If either + invariant broke, the gathered copy would carry a stale scale on + iter-2+ and silently apply the wrong dequantization. + + This test locks the invariant down by forcing a radically different + scale between iterations and asserting the gathered tensor's + dequantization tracks the sharded one. + """ + torch.manual_seed(42) + hybrid_recipe = _hybrid_custom_recipe( + _fp8_row_factory, + _fp8_col_factory, + _fp8_grad_factory, + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + hybrid_param = model.weight + + # Iter-1 gather with the initial (small-magnitude) weights + sharded_tensors_1, metadata_1 = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + gathered, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors_1, + metadata_1, + hybrid_param.dtype, + out=None, + ) + + # Simulate an optimizer writeback that produces a much larger weight; + # Float8CurrentScalingQuantizer.update_quantized must recompute + # _scale_inv for this range. If the gathered copy didn't see the new + # scale, the dequantize below would disagree with the sharded copy. + huge_master = torch.randn_like(hybrid_param.dequantize()) * 100.0 + hybrid_param._quantizer.update_quantized(huge_master, hybrid_param) + + # Iter-2+ path: reuse the gathered buffer + sharded_tensors_2, metadata_2 = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + gathered_refreshed, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors_2, + metadata_2, + hybrid_param.dtype, + out=gathered, + ) + assert gathered_refreshed is gathered + + # The gathered copy must now reflect the new sharded scale, not the + # tiny original scale. + torch.testing.assert_close( + hybrid_param.dequantize(), + gathered_refreshed.dequantize(), + rtol=0.0, + atol=0.0, + ) + _assert_hybrid_tensor_exact(gathered_refreshed, hybrid_param, context="FSDP scale refresh") + # And the magnitude really did change (sanity: this test would pass + # vacuously if update_quantized didn't actually change anything). + assert gathered_refreshed.dequantize().abs().max() > 10.0, ( + "update_quantized did not produce a sufficiently different " + "weight; the scale-refresh invariant is not being exercised" + ) + + def test_nvfp4_sub_storage_raises_on_pre_all_gather(self): + """Hybrid FSDP2 with an NVFP4 sub-storage must raise a clear error. + + NVIDIA/TransformerEngine#3158 tracks the missing NVFP4 FSDP2 support: + packed FP4 dim-0 alignment, columnwise dequantization, and RHT-cache + handling are not implemented. Until that lands, hybrid pre-all-gather must + refuse an NVFP4 sub-storage cleanly via the ``fsdp_buffer_fields`` + protocol rather than silently producing wrong data. + + This test pins that contract: any hybrid whose sub-storage does not + implement ``fsdp_buffer_fields`` raises ``NotImplementedError`` at + ``fsdp_pre_all_gather`` time. The prior version of this test + inadvertently asserted the opposite when buffer extraction used + implicit ``get_metadata()``-based tensor scanning. + """ + if not (fp8_available and nvfp4_available): + pytest.skip("Requires FP8 + NVFP4 support") + + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: NVFP4Quantizer(), + col_factory=lambda: NVFP4Quantizer(), + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + param = model.weight + + # Clean refusal: hybrid's pre_all_gather raises an NVFP4-specific + # message identifying the unsupported sub-storage protocol, not a generic + # "NVFP4Tensor does not implement fsdp_buffer_fields" from deep inside + # the base class. + with pytest.raises(NotImplementedError) as exc_info: + param.fsdp_pre_all_gather( + mesh=None, + orig_size=param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + msg = str(exc_info.value) + assert "NVFP4Tensor" in msg + assert "rowwise sub-storage" in msg + assert "Use a supported sub-quantizer" in msg + assert "fsdp_buffer_fields" in msg + + +# --------------------------------------------------------------------------- +# 5. make_like correctness +# --------------------------------------------------------------------------- + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridMakeLike: + """Test that make_like produces correct copies for __torch_dispatch__ usage.""" + + def _make_hybrid_param(self): + hybrid_recipe = _hybrid_custom_recipe( + _fp8_row_factory, + _fp8_col_factory, + _fp8_grad_factory, + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + return model.weight + + def test_make_like_preserves_sub_storages(self): + """make_like result has the same sub-storage types, quantizers, and dtype.""" + param = self._make_hybrid_param() + copy = HybridQuantizedTensor.make_like(param) + + assert isinstance(copy, HybridQuantizedTensor) + assert copy.dtype == param.dtype + assert copy.shape == param.shape + assert type(copy.rowwise_sub_storage) is type(param.rowwise_sub_storage) + assert type(copy.columnwise_sub_storage) is type(param.columnwise_sub_storage) + _assert_hybrid_tensor_exact(copy, param, context="make_like") + torch.testing.assert_close(copy.dequantize(), param.dequantize(), rtol=0.0, atol=0.0) + + def test_make_like_is_independent(self): + """make_like result should not share the same tensor identity.""" + param = self._make_hybrid_param() + copy = HybridQuantizedTensor.make_like(param) + assert copy is not param + + +# --------------------------------------------------------------------------- +# 5b. Hopper-only paths: columnwise-only Float8 sub-storage +# --------------------------------------------------------------------------- +# +# On architectures where ``is_non_tn_fp8_gemm_supported()`` returns False +# (Hopper sm_90, L40 sm_89), per-tensor FP8 GEMM only supports the TN +# layout — non-TN layouts are simulated by feeding pre-transposed data. +# So a columnwise-only ``Float8TensorStorage`` (used as a hybrid sub- +# storage) holds its quantized data in ``_transpose`` instead of +# ``_data``, with ``_data = None``. +# +# This is the exact layout the FSDP2 buffer protocol must recognize +# when the sub-storage is part of a ``HybridQuantizedTensor`` parameter. +# These tests pin the contracts that would break if the buffer +# protocol regressed to the unconditional ``("_data",)`` field name +# (which would all-gather a ``None`` tensor on Hopper). +# +# Skip on Blackwell where the C++ kernel always populates ``_data`` and +# the columnwise-only Float8 path doesn't exercise ``_transpose``. + +requires_hopper_fp8 = pytest.mark.skipif( + is_non_tn_fp8_gemm_supported() or not fp8_available, + reason=( + "Hopper-only: requires per-tensor FP8 with non-TN GEMM unsupported " + "(Hopper sm_90 / L40 sm_89). On Blackwell the C++ kernel populates " + "_data even for columnwise-only mode, so the _transpose-only path " + "is not exercised." + ), +) + + +@requires_hopper_fp8 +class TestHybridFloat8ColumnwiseOnlyHopperPath: + """Hopper transpose-only Float8 uses an M-major FSDP transport layout.""" + + @staticmethod + def _make_columnwise_only_float8_storage(shape=(32, 64), value=1.0, quantizer=None): + source = torch.full(shape, value, device="cuda", dtype=torch.bfloat16) + byte_quantizer = Float8Quantizer( + scale=torch.ones(1, device="cuda", dtype=torch.float32), + amax=torch.zeros(1, device="cuda", dtype=torch.float32), + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + ) + rowwise = byte_quantizer(source) + if quantizer is None: + quantizer = Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + rowwise=False, + columnwise=True, + ) + storage = Float8Tensor( + shape=shape, + dtype=source.dtype, + data=None, + data_transpose=rowwise._data.movedim(-1, 0).contiguous(), + fp8_scale_inv=rowwise._scale_inv.clone(), + fp8_dtype=tex.DType.kFloat8E4M3, + quantizer=quantizer, + requires_grad=False, + device="cuda", + ) + return storage, rowwise.dequantize() + + @classmethod + def _make_hybrid_shard(cls, value): + quantizer = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + ), + ) + source = torch.full((32, 64), value, device="cuda", dtype=torch.bfloat16) + rowwise = quantizer.rowwise_quantizer(source) + columnwise, expected = cls._make_columnwise_only_float8_storage( + value=value, quantizer=quantizer.columnwise_quantizer + ) + hybrid = HybridQuantizedTensor( + shape=source.shape, + dtype=source.dtype, + rowwise_storage=rowwise, + columnwise_storage=columnwise, + quantizer=quantizer, + requires_grad=False, + device="cuda", + ) + return hybrid, expected + + def test_columnwise_only_float8_fsdp_buffer_fields_returns_transpose(self): + out, _ = self._make_columnwise_only_float8_storage() + + assert out._data is None + assert out._transpose is not None + assert not out._transpose_invalid + assert out.fsdp_buffer_fields() == ("_transpose",) + + def test_columnwise_only_float8_fsdp_extract_uses_m_major_transport(self): + out, _ = self._make_columnwise_only_float8_storage() + + buffers, metadata = out.fsdp_extract_buffers() + + assert len(buffers) == 1 + assert buffers[0].shape == out.shape + assert buffers[0].is_contiguous() + torch.testing.assert_close(buffers[0], out._transpose.movedim(0, -1).contiguous()) + assert metadata == { + "field_names": ("_transpose",), + "transport_layout": "columnwise_m_major", + } + + def test_columnwise_only_float8_fsdp_assign_restores_columnwise_layout(self): + out, expected = self._make_columnwise_only_float8_storage() + buffers, metadata = out.fsdp_extract_buffers() + gathered = torch.cat((buffers[0], buffers[0]), dim=0) + rebuilt = Float8Tensor.make_like(out, shape=gathered.shape) + + rebuilt.fsdp_assign_gathered((gathered,), metadata) + + assert rebuilt.shape == torch.Size((64, 64)) + assert rebuilt._data is None + assert rebuilt._transpose.shape == torch.Size((64, 64)) + assert not rebuilt._transpose_invalid + torch.testing.assert_close(rebuilt._transpose, gathered.movedim(-1, 0).contiguous()) + torch.testing.assert_close(rebuilt.dequantize(), torch.cat((expected, expected), dim=0)) + + def test_hybrid_fsdp_two_rank_gather_and_buffer_reuse(self): + shards = [self._make_hybrid_shard(value) for value in (1.0, 2.0)] + extracted = [ + shard.fsdp_pre_all_gather( + mesh=None, + orig_size=shard.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + for shard, _ in shards + ] + gathered = tuple( + torch.cat((extracted[0][0][i], extracted[1][0][i]), dim=0) + for i in range(len(extracted[0][0])) + ) + expected = torch.cat((shards[0][1], shards[1][1]), dim=0) + + rebuilt, _ = shards[0][0].fsdp_post_all_gather(gathered, extracted[0][1], torch.bfloat16) + + assert rebuilt.shape == torch.Size((64, 64)) + assert rebuilt._columnwise_storage.shape == torch.Size((64, 64)) + assert rebuilt._columnwise_storage._data is None + assert rebuilt._columnwise_storage._transpose.shape == torch.Size((64, 64)) + torch.testing.assert_close(rebuilt._columnwise_storage.dequantize(), expected) + + reused, _ = shards[0][0].fsdp_post_all_gather( + gathered, extracted[0][1], torch.bfloat16, out=rebuilt + ) + assert reused is rebuilt + assert reused._columnwise_storage._transpose.shape == torch.Size((64, 64)) + torch.testing.assert_close(reused._columnwise_storage.dequantize(), expected) + + def test_fsdp_buffer_fields_falls_back_to_data_when_both_present(self): + quantizer = Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + ) + out = quantizer(torch.randn(32, 64, device="cuda", dtype=torch.bfloat16)) + + assert out._data is not None + assert out._transpose is not None + assert out.fsdp_buffer_fields() == ("_data",) + buffers, metadata = out.fsdp_extract_buffers() + assert len(buffers) == 1 + assert buffers[0] is out._data + assert metadata == { + "field_names": ("_data",), + "transport_layout": "native", + } + + +@requires_hopper_fp8 +class TestPerTensorFloat8ColumnwiseOnlyGuard: + """Per-tensor FP8 cannot yet emit only the physical transpose.""" + + @staticmethod + def _make_quantizer(scaling): + if scaling == "current": + return Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + rowwise=False, + columnwise=True, + ) + return Float8Quantizer( + scale=torch.ones(1, device="cuda", dtype=torch.float32), + amax=torch.zeros(1, device="cuda", dtype=torch.float32), + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=False, + columnwise=True, + ) + + @pytest.mark.parametrize("scaling", ("current", "delayed")) + def test_initial_quantize_raises_not_implemented(self, scaling): + quantizer = self._make_quantizer(scaling) + source = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16) + + with pytest.raises( + NotImplementedError, + match="Columnwise-only per-tensor FP8 quantization is not implemented", + ): + quantizer(source) + + @pytest.mark.parametrize("scaling", ("current", "delayed")) + def test_update_quantized_raises_not_implemented(self, scaling): + quantizer = self._make_quantizer(scaling) + source = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16) + output = quantizer.make_empty( + source.shape, + dtype=source.dtype, + device=source.device, + ) + assert output._data is None + assert output._transpose is not None + + with pytest.raises( + NotImplementedError, + match="Columnwise-only per-tensor FP8 quantization is not implemented", + ): + quantizer.update_quantized(source, output) + + +@requires_hopper_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridFsdpPostAllGatherUpdateUsage: + """``HybridQuantizedTensor.fsdp_post_all_gather`` must call + ``update_usage`` on each sub-storage after writing gathered data + (mirroring vanilla ``Float8Tensor.fsdp_post_all_gather:888``). + Without it, on Hopper a previously-cached ``_transpose`` from the + prior iteration is silently reused with the new ``_data``, producing + incorrect dgrad / wgrad GEMMs. + """ + + def _make_param(self): + hybrid_recipe = _hybrid_custom_recipe( + _fp8_row_factory, + _fp8_col_factory, + _fp8_grad_factory, + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(64, 64, params_dtype=torch.bfloat16).cuda() + return model.weight + + def test_iter2_invalidates_stale_transpose_on_rowwise_substorage(self): + """Simulates iter-2+ buffer reuse: pre-existing ``out`` with a + possibly-stale ``_transpose`` cache; after ``fsdp_post_all_gather`` + the rowwise sub-storage's ``_transpose`` must be invalidated / + regenerated to match the freshly gathered ``_data``. + """ + param = self._make_param() + # Build a plausible iter-2+ "out" with stale state. + out = HybridQuantizedTensor.make_like(param) + # Rowwise sub-storage on Hopper has _data populated. Force a stale + # _transpose and invalidate flag to mimic the regression scenario. + if out._rowwise_storage._transpose is None: + # Set up a fake stale transpose (non-None, marked invalid by + # the mismatching shape would catch nothing, so just plant + # a tensor and clear the invalid flag to "valid"). + out._rowwise_storage._transpose = torch.zeros_like(out._rowwise_storage._data).t() + out._rowwise_storage._transpose_invalid = False + stale_transpose_id = id(out._rowwise_storage._transpose) + + # Drive a real all-gather round trip via the protocol + sharded_tensors, metadata = param.fsdp_pre_all_gather( + mesh=None, + orig_size=param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + out2, _ = param.fsdp_post_all_gather(sharded_tensors, metadata, param.dtype, out=out) + + # After fsdp_post_all_gather, the rowwise sub-quantizer is pinned + # columnwise=False, so update_usage(rowwise=True, columnwise=False) + # must clear the stale _transpose (preventing the silent + # stale-cache regression on Hopper). + assert out2._rowwise_storage._transpose is None or ( + out2._rowwise_storage._transpose_invalid + and id(out2._rowwise_storage._transpose) != stale_transpose_id + ), "Stale _transpose was not invalidated after fsdp_post_all_gather" diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py new file mode 100644 index 0000000000..cb0785e3e0 --- /dev/null +++ b/tests/pytorch/test_identity_quantizer.py @@ -0,0 +1,1867 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for IdentityQuantizer (high-precision passthrough) and its use as a +per-direction component of HybridQuantizer to express mixed forward/backward +precision via the CustomRecipe + qfactory machinery. Scoped to single-GPU +TE GEMM modules. +""" + +import io + +import pytest +import torch + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine.common.recipe import CustomRecipe +from transformer_engine.pytorch import ( + Float8BlockQuantizer, + Float8CurrentScalingQuantizer, + HybridQuantizer, + HybridQuantizedTensor, + IdentityQuantizer, + MXFP8Quantizer, + NVFP4Quantizer, +) +from transformer_engine.pytorch.tensor.identity_tensor import IdentityTensor +from transformer_engine.pytorch.tensor.storage.identity_tensor_storage import ( + IdentityTensorStorage, +) +from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported + +fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) +mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) +nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) +fp8_block_scaling_available, reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available( + return_reason=True +) + +_COLUMNWISE_ONLY_PER_TENSOR_FP8_ERROR = ( + "Columnwise-only per-tensor FP8 quantization is not implemented" +) + +_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 = pytest.mark.xfail( + condition=not is_non_tn_fp8_gemm_supported(), + raises=NotImplementedError, + strict=True, + reason=( + "Hopper does not yet support columnwise-only per-tensor FP8 quantization; " + "tracked by NVIDIA/TransformerEngine#3158" + ), +) + + +# ── Module-level qfactories (picklable / autocast-friendly) ────────── + + +def identity_all_factory(role): # pylint: disable=unused-argument + """Whole layer in high precision: Identity for every slot.""" + return IdentityQuantizer() + + +def _fp8_cs(fp8_dtype=tex.DType.kFloat8E4M3): + return Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device="cuda") + + +def _mxfp8(fp8_dtype=tex.DType.kFloat8E4M3): + return MXFP8Quantizer(fp8_dtype=fp8_dtype) + + +def _float8_blockwise(fp8_dtype=tex.DType.kFloat8E4M3): + return Float8BlockQuantizer(fp8_dtype=fp8_dtype, rowwise=True, columnwise=True) + + +def _nvfp4(): + return NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) + + +_HYBRID_IDENTITY_FORMATS = [ + pytest.param( + "fp8_current", + marks=pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}"), + ), + pytest.param( + "mxfp8", + marks=pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}"), + ), + pytest.param( + "float8_blockwise", + marks=pytest.mark.skipif( + not fp8_block_scaling_available, + reason=f"Float8Blockwise: {reason_for_no_fp8_block_scaling}", + ), + ), + pytest.param( + "nvfp4", + marks=pytest.mark.skipif( + not (fp8_available and nvfp4_available), + reason=f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", + ), + ), +] + +_HYBRID_IDENTITY_RECOMPUTE_FORMATS = [ + pytest.param( + "fp8_current", + marks=pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}"), + ), + pytest.param( + "mxfp8", + marks=pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}"), + ), + pytest.param( + "nvfp4", + marks=pytest.mark.skipif( + not (fp8_available and nvfp4_available), + reason=f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", + ), + ), +] + + +def _format_quantizer(format_name): + if format_name == "fp8_current": + return _fp8_cs(tex.DType.kFloat8E4M3) + if format_name == "mxfp8": + return _mxfp8(tex.DType.kFloat8E4M3) + if format_name == "float8_blockwise": + return _float8_blockwise(tex.DType.kFloat8E4M3) + if format_name == "nvfp4": + return _nvfp4() + raise ValueError(format_name) + + +def _hybrid_quantized_fwd_identity_bwd_factory(format_name): + def qfactory(role): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return IdentityQuantizer() + return HybridQuantizer( + rowwise_quantizer=_format_quantizer(format_name), + columnwise_quantizer=IdentityQuantizer(), + ) + + return qfactory + + +def fwd_hp_bwd_fp8_factory(role): + """High-precision forward, FP8 backward (per-direction via hybrid).""" + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return _fp8_cs(tex.DType.kFloat8E5M2) + return HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=_fp8_cs(tex.DType.kFloat8E4M3), + ) + + +def fwd_fp8_bwd_hp_factory(role): + """FP8 forward, high-precision backward (per-direction via hybrid).""" + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return IdentityQuantizer() + return HybridQuantizer( + rowwise_quantizer=_fp8_cs(tex.DType.kFloat8E4M3), + columnwise_quantizer=IdentityQuantizer(), + ) + + +def fwd_fp8_bwd_rowwise_dequantized_hp_factory(role): + """FP8 forward, high-precision backward from dequantized fprop values.""" + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return IdentityQuantizer() + return HybridQuantizer( + rowwise_quantizer=_fp8_cs(tex.DType.kFloat8E4M3), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + + +def hybrid_all_identity_factory(role): + """All directions high precision, expressed through the hybrid container. + + weight / input / output -> Hybrid(Identity, Identity); grad -> Identity. + Exercises the HybridQuantizedTensor path with Identity sub-storages in both + directions (distinct from the non-hybrid whole-layer-HP path). + """ + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return IdentityQuantizer() + return HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + + +def fp8_fwd_factory(role): + """Plain FP8 current scaling for every slot (E4M3 fwd, E5M2 grad). + + Used with ``backward_override="high_precision"`` as the reference that the + per-direction Identity machinery (``fwd_fp8_bwd_hp_factory``) must reproduce + bitwise: same FP8 forward, high-precision backward. + """ + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return _fp8_cs(tex.DType.kFloat8E5M2) + return _fp8_cs(tex.DType.kFloat8E4M3) + + +def _offload_roundtrip(tensor): + from transformer_engine.pytorch.cpu_offload import OffloadableLayerState + + stream = torch.cuda.Stream() + state = OffloadableLayerState(offload_stream=stream) + tid = state.push_tensor(tensor) + state.start_offload() + state.release_activation_forward_gpu_memory() + state.start_reload() + reloaded = state.pop_tensor(tid) + torch.cuda.synchronize() + try: + return reloaded + finally: + state.release_all_memory() + + +# ── Unit tests ─────────────────────────────────────────────────────── + + +class TestIdentityQuantizerUnit: + """IdentityQuantizer / IdentityTensorStorage basic behavior.""" + + def test_quantize_returns_identity_tensor(self): + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + out = IdentityQuantizer()(x) + assert isinstance(out, IdentityTensor) + + def test_internal_returns_storage(self): + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + q = IdentityQuantizer() + q.internal = True + out = q(x) + assert isinstance(out, IdentityTensorStorage) + assert not isinstance(out, IdentityTensor) + + def test_make_empty_internal_returns_storage(self): + q = IdentityQuantizer() + q.internal = True + + out = q.make_empty((8, 16), dtype=torch.bfloat16, device="cuda") + + assert isinstance(out, IdentityTensorStorage) + assert not isinstance(out, IdentityTensor) + assert out.size() == torch.Size((8, 16)) + assert out.dequantize().dtype == torch.bfloat16 + + def test_grouped_split_all_identity_uses_plain_tensor_views(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _split_quantize_non_hybrid, + ) + + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + m_splits = [3, 5] + quantizers = [IdentityQuantizer(), IdentityQuantizer()] + + out = _split_quantize_non_hybrid(x, m_splits, quantizers, activation_dtype=torch.bfloat16) + + assert all(isinstance(t, torch.Tensor) for t in out) + assert not any(isinstance(t, IdentityTensorStorage) for t in out) + for actual, expected in zip(out, torch.split(x, m_splits)): + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + + cast_quantizers = [ + IdentityQuantizer(dtype=torch.float32), + IdentityQuantizer(dtype=torch.float32), + ] + cast_out = _split_quantize_non_hybrid( + x, m_splits, cast_quantizers, activation_dtype=torch.bfloat16 + ) + assert all(isinstance(t, IdentityTensorStorage) for t in cast_out) + for actual, expected in zip(cast_out, torch.split(x, m_splits)): + dequantized = actual.dequantize() + assert dequantized.dtype == torch.float32 + torch.testing.assert_close( + dequantized, + expected.to(torch.float32), + rtol=0.0, + atol=0.0, + ) + + def test_grouped_split_rejects_mixed_identity_and_quantized_operands(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _validate_grouped_quantizer_list, + ) + + cases = [ + [IdentityQuantizer(), _mxfp8(tex.DType.kFloat8E4M3)], + [ + HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ), + HybridQuantizer( + rowwise_quantizer=_mxfp8(tex.DType.kFloat8E4M3), + columnwise_quantizer=IdentityQuantizer(), + ), + ], + ] + + for quantizers in cases: + with pytest.raises(ValueError, match="mix Identity-backed and quantized"): + _validate_grouped_quantizer_list(quantizers, operand_name="input") + + def test_hybrid_split_forwards_disable_bulk_allocation_to_both_directions(self, monkeypatch): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + from transformer_engine.pytorch.module.grouped_linear import _split_quantize_hybrid + + calls = [] + + def fake_split_quantize(tensor, m_splits, quantizers, *, disable_bulk_allocation=False): + calls.append(disable_bulk_allocation) + return [ + quantizer(tensor_part) + for tensor_part, quantizer in zip(torch.split(tensor, m_splits), quantizers) + ] + + monkeypatch.setattr(grouped_linear.tex, "split_quantize", fake_split_quantize) + monkeypatch.setattr( + grouped_linear, + "_supports_native_split_quantize", + lambda quantizer: True, + ) + x = torch.randn(8, 16, dtype=torch.bfloat16) + m_splits = [3, 5] + quantizers = [ + HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + for _ in m_splits + ] + + out = _split_quantize_hybrid( + x, + m_splits, + quantizers, + disable_bulk_allocation=True, + ) + + assert calls == [True, True] + assert len(out) == len(m_splits) + + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + def test_grouped_linear_cpu_offload_disables_bulk_allocation_for_hybrid_input( + self, monkeypatch + ): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + class StopAfterFlagCapture(RuntimeError): + pass + + def qfactory(role): + if role is not None and role.module_type == "grouped_linear": + return HybridQuantizer( + rowwise_quantizer=_fp8_cs(tex.DType.kFloat8E4M3), + columnwise_quantizer=_fp8_cs(tex.DType.kFloat8E4M3), + ) + return _fp8_cs(tex.DType.kFloat8E4M3) + + calls = [] + + def fake_split_quantize_hybrid( + tensor, m_splits, quantizers, *, disable_bulk_allocation=False, **kwargs + ): + del tensor, m_splits, quantizers, kwargs + calls.append(disable_bulk_allocation) + raise StopAfterFlagCapture("captured hybrid split kwargs") + + monkeypatch.setattr(grouped_linear, "is_cpu_offload_enabled", lambda: True) + monkeypatch.setattr(grouped_linear, "_split_quantize_hybrid", fake_split_quantize_hybrid) + + model = te.GroupedLinear(2, 64, 64, params_dtype=torch.bfloat16).cuda() + x = torch.randn(64, 64, device="cuda", dtype=torch.bfloat16) + m_splits = torch.tensor([32, 32], device="cuda", dtype=torch.int32) + + with pytest.raises(StopAfterFlagCapture): + with te.autocast(enabled=True, recipe=CustomRecipe(qfactory=qfactory)): + model(x, m_splits=m_splits) + + assert calls == [True] + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_grouped_linear_rejects_mixed_identity_weight_quantizers(self): + weight_count = 0 + + def qfactory(role): + nonlocal weight_count + if role is not None and role.module_type == "grouped_linear": + if role.tensor_type == "weight": + weight_count += 1 + if weight_count == 1: + return IdentityQuantizer() + return _mxfp8(tex.DType.kFloat8E4M3) + return _mxfp8(tex.DType.kFloat8E4M3) + + model = te.GroupedLinear(2, 64, 64, params_dtype=torch.bfloat16).cuda() + x = torch.randn(64, 64, device="cuda", dtype=torch.bfloat16) + m_splits = torch.tensor([32, 32], device="cuda", dtype=torch.int32) + + with pytest.raises(ValueError, match="mix Identity-backed and quantized"): + with te.autocast(enabled=True, recipe=CustomRecipe(qfactory=qfactory)): + model(x, m_splits=m_splits) + + def test_identity_contiguous_preserves_wrapper_and_values(self): + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16).t() + t = IdentityQuantizer()(x) + + out = t.contiguous() + + assert isinstance(out, IdentityTensor) + assert out.is_contiguous() + torch.testing.assert_close(out.dequantize(), x.contiguous(), rtol=0.0, atol=0.0) + + def test_hybrid_identity_contiguous_preserves_wrapper_and_values(self): + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + q = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + t = q(x) + + out = t.contiguous() + + assert out is t + assert isinstance(out, HybridQuantizedTensor) + assert isinstance(out.rowwise_sub_storage, IdentityTensor) + assert isinstance(out.columnwise_sub_storage, IdentityTensor) + torch.testing.assert_close(out.dequantize(), x, rtol=0.0, atol=0.0) + + def test_hybrid_identity_cpu_preserves_nested_storage_types(self): + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + q = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + t = q(x) + + out = t.cpu() + + assert isinstance(out, HybridQuantizedTensor) + assert out.device.type == "cpu" + assert isinstance(out.rowwise_sub_storage, IdentityTensor) + assert isinstance(out.columnwise_sub_storage, IdentityTensor) + assert out.rowwise_sub_storage.device.type == "cpu" + assert out.columnwise_sub_storage.device.type == "cpu" + torch.testing.assert_close(out.dequantize(), x.cpu(), rtol=0.0, atol=0.0) + assert len(out.get_data_tensors()) == 4 + out.copy_(torch.ones_like(x, device="cpu")) + torch.testing.assert_close( + out.dequantize(), torch.ones_like(x, device="cpu"), rtol=0.0, atol=0.0 + ) + + def test_hybrid_quantizer_copy_preserves_parent_flags(self): + q = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + q.set_usage(rowwise=True, columnwise=False) + q.internal = True + q.optimize_for_gemm = True + + out = q.copy() + + assert isinstance(out, HybridQuantizer) + assert out is not q + assert out.rowwise_quantizer is not q.rowwise_quantizer + assert out.columnwise_quantizer is not q.columnwise_quantizer + assert out.rowwise_usage is True + assert out.columnwise_usage is False + assert out.internal is True + assert out.optimize_for_gemm is True + assert out.rowwise_quantizer.rowwise_usage is True + assert out.rowwise_quantizer.columnwise_usage is False + assert out.columnwise_quantizer.rowwise_usage is False + assert out.columnwise_quantizer.columnwise_usage is True + + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + def test_te_ops_basic_linear_accepts_hybrid_identity_quantized_weight(self): + import transformer_engine.pytorch.ops as te_ops + + def qfactory(role): # pylint: disable=unused-argument + return HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + + torch.manual_seed(1701) + ref = te_ops.BasicLinear(16, 16, device="cuda", dtype=torch.bfloat16) + custom_recipe = CustomRecipe(qfactory=qfactory) + torch.manual_seed(1702) + with te.quantized_model_init(enabled=True, recipe=custom_recipe): + test = te_ops.BasicLinear(16, 16, device="cuda", dtype=torch.bfloat16) + + with torch.no_grad(): + test.weight.copy_(ref.weight) + + torch.manual_seed(1703) + x = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16) + x_ref = x.detach().clone().requires_grad_(True) + x_test = x.detach().clone().requires_grad_(True) + + y_ref = ref(x_ref) + with te.autocast(enabled=True, recipe=custom_recipe): + y_test = test(x_test) + + torch.manual_seed(1704) + grad_output = torch.randn_like(y_ref) + y_ref.backward(grad_output) + y_test.backward(grad_output) + + assert isinstance(test.weight, HybridQuantizedTensor) + torch.testing.assert_close(y_test, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(x_test.grad, x_ref.grad, rtol=0.0, atol=0.0) + torch.testing.assert_close(test.weight.grad, ref.weight.grad, rtol=0.0, atol=0.0) + + @pytest.mark.parametrize( + "qfactory", + [ + pytest.param(lambda role: IdentityQuantizer(), id="identity"), + pytest.param( + lambda role: HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ), + id="hybrid_identity", + ), + ], + ) + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + def test_te_ops_quantize_then_gelu_accepts_identity_backed_tensors(self, qfactory): + import transformer_engine.pytorch.ops as te_ops + + ref = te_ops.GELU() + test = te_ops.Sequential(te_ops.Quantize(forward=True), te_ops.GELU()) + x = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16) + x_ref = x.detach().clone().requires_grad_(True) + x_test = x.detach().clone().requires_grad_(True) + + y_ref = ref(x_ref) + with te.autocast(enabled=True, recipe=CustomRecipe(qfactory=qfactory)): + y_test = test(x_test) + + torch.manual_seed(1801) + grad_output = torch.randn_like(y_ref) + y_ref.backward(grad_output) + y_test.backward(grad_output) + + torch.testing.assert_close(y_test, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(x_test.grad, x_ref.grad, rtol=0.0, atol=0.0) + + def test_hybrid_fsdp_rejects_storage_only_sub_storages(self): + row_quantizer = IdentityQuantizer() + col_quantizer = IdentityQuantizer() + row_quantizer.internal = True + col_quantizer.internal = True + q = HybridQuantizer( + rowwise_quantizer=row_quantizer, + columnwise_quantizer=col_quantizer, + ) + t = q(torch.randn(8, 16, device="cuda", dtype=torch.bfloat16)) + + with pytest.raises(NotImplementedError, match="storage-only rowwise sub-storage"): + t.fsdp_pre_all_gather( + mesh=None, + orig_size=t.shape, + contiguous_orig_stride=t.stride(), + module=None, + mp_policy=None, + ) + + def test_hybrid_quantizer_rejects_nested_quantizer_requests(self): + from transformer_engine.pytorch.quantization import DelayedScalingRequest + + with pytest.raises(TypeError, match="does not support nested QuantizerRequest"): + HybridQuantizer( + rowwise_quantizer=DelayedScalingRequest(), + columnwise_quantizer=IdentityQuantizer(), + ) + + def test_fp8_dpa_rejects_identity_quantizer_with_type_error(self): + from transformer_engine.pytorch.attention.dot_product_attention import utils as dpa_utils + from transformer_engine.pytorch.cpp_extensions.fused_attn import ( + META_DO, + META_DP, + META_DQKV, + META_O, + META_QKV, + META_S, + ) + + n_fwd = max(META_QKV, META_S, META_O) + 1 + n_bwd = max(META_DO, META_DP, META_DQKV) + 1 + quantizers = { + "scaling_fwd": [IdentityQuantizer() for _ in range(n_fwd)], + "scaling_bwd": [IdentityQuantizer() for _ in range(n_bwd)], + } + + with pytest.raises(TypeError, match="FP8 attention requires FP8-compatible quantizers"): + dpa_utils.get_attention_quantizers(True, quantizers) + + def test_dequantize_bitwise_identical(self): + x = torch.randn(4, 32, device="cuda", dtype=torch.bfloat16) + out = IdentityQuantizer()(x) + assert torch.equal(out.dequantize(), x) + + def test_dtype_cast(self): + x = torch.randn(4, 8, device="cuda", dtype=torch.float32) + out = IdentityQuantizer(dtype=torch.bfloat16)(x) + dequantized = out.dequantize() + assert dequantized.dtype == torch.bfloat16 + torch.testing.assert_close(dequantized, x.to(torch.bfloat16), rtol=0.0, atol=0.0) + + def test_make_empty_honors_configured_dtype(self): + tensor = IdentityQuantizer(dtype=torch.float32).make_empty( + (4, 8), + dtype=torch.bfloat16, + device="cuda", + ) + + assert tensor.dtype == torch.float32 + assert tensor._hp_data.dtype == torch.float32 + assert tensor.dequantize().dtype == torch.float32 + + def test_update_quantized_synchronizes_dtype(self): + dst = IdentityQuantizer().make_empty( + (4, 8), + dtype=torch.bfloat16, + device="cuda", + ) + quantizer = IdentityQuantizer(dtype=torch.float32) + + quantizer.update_quantized(torch.ones_like(dst._hp_data), dst) + + assert dst.dtype == torch.float32 + assert dst._hp_data.dtype == torch.float32 + assert dst.dequantize().dtype == torch.float32 + + @pytest.mark.parametrize("noop", [0, 1]) + def test_update_quantized_honors_noop(self, noop): + quantizer = IdentityQuantizer() + dst = quantizer(torch.full((4, 8), 3.0, device="cuda")) + src = torch.full((4, 8), 7.0, device="cuda") + + quantizer.update_quantized( + src, dst, noop_flag=torch.tensor(noop, dtype=torch.float32, device="cuda") + ) + + expected = 3.0 if noop else 7.0 + torch.testing.assert_close(dst.dequantize(), torch.full_like(src, expected)) + + @pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") + def test_hybrid_update_honors_shared_noop(self): + quantizer = HybridQuantizer( + rowwise_quantizer=_fp8_cs(), + columnwise_quantizer=IdentityQuantizer(), + ) + old = torch.full((32, 32), 3.0, dtype=torch.bfloat16, device="cuda") + dst = quantizer(old) + expected_row = dst._rowwise_storage.dequantize().clone() + + quantizer.update_quantized( + torch.full_like(old, 7.0), + dst, + noop_flag=torch.ones(1, dtype=torch.float32, device="cuda"), + ) + + torch.testing.assert_close(dst._rowwise_storage.dequantize(), expected_row) + torch.testing.assert_close(dst._columnwise_storage.dequantize(), old) + + def test_hybrid_new_zeros_preserves_identity_dtypes(self): + quantizer = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(dtype=torch.bfloat16), + columnwise_quantizer=IdentityQuantizer(dtype=torch.float32), + ) + tensor = quantizer(torch.ones(4, 8, dtype=torch.bfloat16, device="cuda")) + + zeros = tensor.new_zeros(tensor.shape) + + assert zeros.rowwise_sub_storage._hp_data.dtype == torch.bfloat16 + assert zeros.columnwise_sub_storage._hp_data.dtype == torch.float32 + + def test_storage_dequantize_defaults_to_nominal_dtype(self): + payload = torch.randn(4, 8, dtype=torch.float32) + storage = IdentityTensorStorage( + hp_data=payload, + fake_dtype=torch.bfloat16, + ) + + default = storage.dequantize() + assert default.dtype == torch.bfloat16 + torch.testing.assert_close(default, payload.to(torch.bfloat16), rtol=0.0, atol=0.0) + + explicit = storage.dequantize(dtype=torch.float16) + assert explicit.dtype == torch.float16 + torch.testing.assert_close(explicit, payload.to(torch.float16), rtol=0.0, atol=0.0) + + same_dtype_storage = IdentityTensorStorage( + hp_data=payload, + fake_dtype=torch.float32, + ) + assert same_dtype_storage.dequantize() is payload + + def test_update_usage_is_noop(self): + x = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) + q = IdentityQuantizer() + q.internal = True + st = q(x) + st.update_usage(rowwise_usage=False, columnwise_usage=True) + assert torch.equal(st.dequantize(), x) + assert st.get_usages() == {"rowwise": True, "columnwise": True} + + def test_save_restore_roundtrip(self): + x = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) + q = IdentityQuantizer() + q.internal = True + st = q(x) + tensors, _ = st.prepare_for_saving() + assert st._hp_data is None + leftover = st.restore_from_saved(tensors) + assert leftover == [] + assert torch.equal(st.dequantize(), x) + + def test_update_quantized_inplace(self): + x = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) + q = IdentityQuantizer() + st = q.make_empty((4, 8), dtype=torch.bfloat16, device="cuda") + q.update_quantized(x, st) + assert torch.equal(st.dequantize(), x) + + def test_tensor_ops_preserve_identity_and_values(self): + x = torch.arange(24, device="cuda", dtype=torch.bfloat16).reshape(6, 4) + t = IdentityQuantizer()(x) + + view = t.view(3, 8) + assert isinstance(view, IdentityTensor) + torch.testing.assert_close(view.dequantize(), x.view(3, 8), rtol=0.0, atol=0.0) + + pieces = torch.split(t, 2, dim=0) + assert all(isinstance(piece, IdentityTensor) for piece in pieces) + for piece, ref in zip(pieces, torch.split(x, 2, dim=0)): + torch.testing.assert_close(piece.dequantize(), ref, rtol=0.0, atol=0.0) + + sliced = t[1:5:2] + assert isinstance(sliced, IdentityTensor) + torch.testing.assert_close(sliced.dequantize(), x[1:5:2], rtol=0.0, atol=0.0) + + strided = torch.as_strided(t, (3, 4), (4, 1), 4) + assert isinstance(strided, IdentityTensor) + torch.testing.assert_close(strided.dequantize(), x[1:4], rtol=0.0, atol=0.0) + + cloned = torch.clone(t) + assert isinstance(cloned, IdentityTensor) + torch.testing.assert_close(cloned.dequantize(), x, rtol=0.0, atol=0.0) + + zeros = t.new_zeros((2, 3)) + assert isinstance(zeros, IdentityTensor) + torch.testing.assert_close( + zeros.dequantize(), + torch.zeros((2, 3), device="cuda", dtype=x.dtype), + rtol=0.0, + atol=0.0, + ) + + dst = IdentityQuantizer().make_empty(x.shape, dtype=x.dtype, device="cuda") + dst.copy_(t) + torch.testing.assert_close(dst.dequantize(), x, rtol=0.0, atol=0.0) + + def test_view_ops_preserve_offsets_strides_and_aliasing(self): + base = torch.arange(24, device="cuda", dtype=torch.float32).reshape(6, 4) + tensor = IdentityQuantizer()(base) + + sliced = torch.ops.aten.slice.Tensor(tensor, 0, 2, 6, 1) + omitted_offset = torch.ops.aten.as_strided.default( + sliced, + [2, 4], + [4, 1], + ) + explicit_offset = torch.ops.aten.as_strided.default( + sliced, + [2, 4], + [4, 1], + 0, + ) + torch.testing.assert_close(omitted_offset.dequantize(), base[2:4], rtol=0.0, atol=0.0) + torch.testing.assert_close(explicit_offset.dequantize(), base[:2], rtol=0.0, atol=0.0) + assert sliced.storage_offset() == base[2:].storage_offset() + assert omitted_offset.storage_offset() == base[2:].storage_offset() + assert explicit_offset.storage_offset() == 0 + + noncontiguous = torch.ops.aten.slice.Tensor(tensor, 1, 0, 4, 2) + assert noncontiguous.stride() == base[:, ::2].stride() + assert noncontiguous.storage_offset() == base[:, ::2].storage_offset() + replacement = torch.full_like(base[:, ::2], -3) + noncontiguous.copy_(replacement) + torch.testing.assert_close(base[:, ::2], replacement, rtol=0.0, atol=0.0) + + pieces = torch.split(tensor, 2, dim=0) + piece_value = torch.full_like(base[:2], 7) + pieces[0].copy_(piece_value) + torch.testing.assert_close(base[:2], piece_value, rtol=0.0, atol=0.0) + + def test_detach_and_clone_preserve_view_metadata(self): + base = torch.arange(24, device="cuda", dtype=torch.float32).reshape(6, 4) + tensor = IdentityQuantizer()(base) + + offset_view = torch.ops.aten.slice.Tensor(tensor, 0, 2, 6, 1) + detached = offset_view.detach() + assert detached.stride() == detached._hp_data.stride() == offset_view.stride() + assert detached.storage_offset() == detached._hp_data.storage_offset() == 8 + replacement = torch.full_like(base[2:6], -5) + detached.copy_(replacement) + torch.testing.assert_close(base[2:6], replacement, rtol=0.0, atol=0.0) + + base2 = torch.arange(24, device="cuda", dtype=torch.float32).reshape(6, 4) + tensor2 = IdentityQuantizer()(base2) + noncontiguous = torch.ops.aten.slice.Tensor(tensor2, 1, 0, 4, 2) + detached_noncontiguous = noncontiguous.detach() + assert detached_noncontiguous.stride() == detached_noncontiguous._hp_data.stride() + assert detached_noncontiguous.storage_offset() == 0 + + transposed = torch.ops.aten.as_strided.default( + tensor2, + [4, 6], + [1, 4], + ) + cloned = transposed.clone() + assert cloned.stride() == cloned._hp_data.stride() + assert cloned.storage_offset() == cloned._hp_data.storage_offset() + torch.testing.assert_close(cloned.dequantize(), base2.as_strided((4, 6), (1, 4))) + original_snapshot = base2.clone() + cloned.copy_(torch.zeros_like(cloned.dequantize())) + torch.testing.assert_close(base2, original_snapshot, rtol=0.0, atol=0.0) + + def test_copy_from_plain_and_identity_sources(self): + dst_data = torch.zeros(3, 4, device="cuda", dtype=torch.float32) + dst = IdentityQuantizer()(dst_data) + + plain_src = torch.arange(12, device="cuda", dtype=torch.float32).reshape(3, 4) + assert dst.copy_(plain_src) is dst + torch.testing.assert_close(dst.dequantize(), plain_src, rtol=0.0, atol=0.0) + + identity_src_data = plain_src.neg() + identity_src = IdentityQuantizer()(identity_src_data) + assert dst.copy_(identity_src) is dst + torch.testing.assert_close(dst.dequantize(), identity_src_data, rtol=0.0, atol=0.0) + + def test_fsdp_pre_post_all_gather_roundtrip(self): + x = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) + t = IdentityQuantizer()(x) + sharded_tensors, metadata = t.fsdp_pre_all_gather( + mesh=None, orig_size=t.shape, contiguous_orig_stride=None, module=None, mp_policy=None + ) + gathered, outputs = t.fsdp_post_all_gather(sharded_tensors, metadata, t.dtype, out=None) + assert isinstance(gathered, IdentityTensor) + assert outputs is sharded_tensors + torch.testing.assert_close(gathered.dequantize(), x, rtol=0.0, atol=0.0) + + reuse, _ = t.fsdp_post_all_gather(sharded_tensors, metadata, t.dtype, out=gathered) + assert reuse is gathered + torch.testing.assert_close(reuse.dequantize(), x, rtol=0.0, atol=0.0) + + def test_torch_weights_only_load_preserves_identity_tensor(self): + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + t = IdentityQuantizer()(x) + buffer = io.BytesIO() + torch.save(t, buffer) + buffer.seek(0) + + loaded = torch.load(buffer, weights_only=True) + + assert isinstance(loaded, IdentityTensor) + assert isinstance(loaded._quantizer, IdentityQuantizer) + torch.testing.assert_close(loaded.dequantize(), x, rtol=0.0, atol=0.0) + + def test_torch_weights_only_load_preserves_hybrid_identity_tensor(self): + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + q = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + t = q(x) + buffer = io.BytesIO() + torch.save(t, buffer) + buffer.seek(0) + + loaded = torch.load(buffer, weights_only=True) + + assert isinstance(loaded, HybridQuantizedTensor) + assert isinstance(loaded._quantizer, HybridQuantizer) + assert isinstance(loaded._rowwise_storage, IdentityTensor) + assert isinstance(loaded._columnwise_storage, IdentityTensor) + torch.testing.assert_close(loaded.dequantize(), x, rtol=0.0, atol=0.0) + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_torch_weights_only_load_preserves_hybrid_mxfp8_identity_tensor(self): + x = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16) + q = HybridQuantizer( + rowwise_quantizer=_mxfp8(tex.DType.kFloat8E4M3), + columnwise_quantizer=IdentityQuantizer(), + ) + t = q(x) + expected = t.dequantize() + buffer = io.BytesIO() + torch.save(t, buffer) + buffer.seek(0) + + loaded = torch.load(buffer, weights_only=True) + + assert isinstance(loaded, HybridQuantizedTensor) + assert isinstance(loaded._quantizer, HybridQuantizer) + assert isinstance(loaded._columnwise_storage, IdentityTensor) + torch.testing.assert_close(loaded.dequantize(), expected, rtol=0.0, atol=0.0) + + def test_cpu_offload_roundtrip_identity_exact(self): + x = torch.randn(1024, 1024, device="cuda", dtype=torch.bfloat16) + t = IdentityQuantizer()(x) + + reloaded = _offload_roundtrip(t) + + assert isinstance(reloaded, IdentityTensor) + torch.testing.assert_close(reloaded.dequantize(), x, rtol=0.0, atol=0.0) + + def test_replace_raw_data_preserves_identity_values(self): + from transformer_engine.pytorch.tensor.utils import replace_raw_data + + x = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) + t = IdentityQuantizer()(x) + new_raw = torch.empty_like(x) + replace_raw_data(t, new_raw) + assert t._hp_data is new_raw + torch.testing.assert_close(t.dequantize(), x, rtol=0.0, atol=0.0) + + def test_quantize_master_weights_identity_exact_nonzero_offset(self): + from transformer_engine.pytorch.tensor.utils import ( + post_all_gather_processing, + quantize_master_weights, + ) + + group = _ensure_single_rank_dp_group() + q = IdentityQuantizer() + weight = q.make_empty((4, 8), dtype=torch.bfloat16, device="cuda") + original = torch.randn_like(weight.dequantize()) + q.update_quantized(original, weight) + + master_full = torch.randn(4, 8, device="cuda", dtype=torch.float32) + start_offset = master_full.numel() // 2 + master_shard = master_full.reshape(-1)[start_offset:].contiguous() + + quantize_master_weights([weight], [master_shard], [start_offset], group=group) + post_all_gather_processing([weight]) + + expected = original.clone() + expected.reshape(-1)[start_offset:] = master_shard.to(torch.bfloat16) + torch.testing.assert_close(weight.dequantize(), expected, rtol=0.0, atol=0.0) + + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_quantize_master_weights_hybrid_identity_fp8_current(self): + from transformer_engine.pytorch.tensor.utils import ( + post_all_gather_processing, + quantize_master_weights, + ) + + group = _ensure_single_rank_dp_group() + recipe = CustomRecipe(qfactory=fwd_hp_bwd_fp8_factory) + torch.manual_seed(123) + with te.quantized_model_init(enabled=True, recipe=recipe): + model = te.Linear(32, 32, bias=False, params_dtype=torch.bfloat16).cuda() + weight = model.weight + assert isinstance(weight, HybridQuantizedTensor) + assert isinstance(weight._rowwise_storage, IdentityTensorStorage) + + master = torch.randn_like(weight.dequantize(dtype=torch.float32)).reshape(-1).contiguous() + quantize_master_weights([weight], [master], [0], group=group) + post_all_gather_processing([weight]) + + expected = master.to(torch.bfloat16) + row_deq = weight._rowwise_storage.dequantize().reshape(-1) + col_deq = weight._columnwise_storage.dequantize(dtype=torch.float32).reshape(-1) + torch.testing.assert_close(row_deq, expected, rtol=0.0, atol=0.0) + torch.testing.assert_close(col_deq, master, rtol=0.125, atol=0.1) + + +# ── te.Linear integration ──────────────────────────────────────────── + + +def _make_linears(in_f, out_f, seed=1234, dtype=torch.bfloat16, bias=True): + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + ref = te.Linear(in_f, out_f, bias=bias, params_dtype=dtype).cuda() + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + test = te.Linear(in_f, out_f, bias=bias, params_dtype=dtype).cuda() + with torch.no_grad(): + for p_test, p_ref in zip(test.parameters(), ref.parameters()): + p_test.copy_(p_ref) + return ref, test + + +def _rel_l2_error(actual, reference): + """Relative L2-norm error ``||actual - reference|| / ||reference||``. + + The right metric for comparing a quantized result to a high-precision + reference: element-wise ``rtol`` is meaningless here because reference grads + contain near-zero entries (relative error on ~1e-12 values explodes), while + the aggregate norm error reflects the true quantization noise. + """ + a = actual.float() + b = reference.float() + return (a - b).norm().item() / (b.norm().item() + 1e-12) + + +def _ensure_single_rank_dp_group(): + import pathlib + import tempfile + + if not torch.distributed.is_initialized(): + torch.cuda.set_device(0) + with tempfile.NamedTemporaryFile(delete=False) as f: + rendezvous_file = pathlib.Path(f.name) + torch.distributed.init_process_group( + backend="nccl", + init_method=rendezvous_file.resolve().as_uri(), + rank=0, + world_size=1, + ) + return torch.distributed.GroupMember.WORLD + + +def _fwd_bwd(model, x, recipe=None): + x = x.clone().detach().requires_grad_(True) + if recipe is not None: + with te.autocast(enabled=True, recipe=recipe): + y = model(x) + else: + y = model(x) + torch.manual_seed(99) + target = torch.randn_like(y) + loss = torch.nn.functional.mse_loss(y, target) + loss.backward() + wgrads = [p.grad.detach().clone() for p in model.parameters() if p.grad is not None] + return y.detach().clone(), x.grad.detach().clone(), wgrads + + +def _fwd_bwd_checkpoint(model, x, recipe, use_reentrant): + x = x.clone().detach().requires_grad_(True) + with te.autocast(enabled=True, recipe=recipe): + if use_reentrant is None: + y = model(x) + else: + y = te.checkpoint(model, x, use_reentrant=use_reentrant) + torch.manual_seed(99) + target = torch.randn_like(y) + loss = torch.nn.functional.mse_loss(y, target) + loss.backward() + wgrads = [p.grad.detach().clone() for p in model.parameters() if p.grad is not None] + return y.detach().clone(), x.grad.detach().clone(), wgrads + + +_IDENTITY_MODULE_NAMES = ( + "Linear", + "LayerNormLinear", + "LayerNormMLP", + "GroupedLinear", + "TransformerLayer", +) + + +def _make_identity_module(module_name, seed=1234, dtype=torch.bfloat16, bias=True): + hidden_size = 64 + ffn_hidden_size = 128 + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + if module_name == "Linear": + return te.Linear(hidden_size, hidden_size, bias=bias, params_dtype=dtype).cuda() + if module_name == "LayerNormLinear": + return te.LayerNormLinear(hidden_size, hidden_size, bias=bias, params_dtype=dtype).cuda() + if module_name == "LayerNormMLP": + return te.LayerNormMLP(hidden_size, ffn_hidden_size, bias=bias, params_dtype=dtype).cuda() + if module_name == "GroupedLinear": + return te.GroupedLinear(2, hidden_size, hidden_size, bias=bias, params_dtype=dtype).cuda() + if module_name == "TransformerLayer": + return te.TransformerLayer( + hidden_size, + ffn_hidden_size, + 4, + hidden_dropout=0.0, + attention_dropout=0.0, + bias=bias, + params_dtype=dtype, + ).cuda() + raise ValueError(module_name) + + +def _make_identity_module_pair(module_name, seed=1234, dtype=torch.bfloat16, bias=True): + ref = _make_identity_module(module_name, seed=seed, dtype=dtype, bias=bias) + test = _make_identity_module(module_name, seed=seed + 1, dtype=dtype, bias=bias) + with torch.no_grad(): + for p_test, p_ref in zip(test.parameters(), ref.parameters()): + p_test.copy_(p_ref) + return ref, test + + +def _identity_module_input(module_name): + torch.manual_seed(7) + if module_name == "TransformerLayer": + return torch.randn(4, 2, 64, device="cuda", dtype=torch.bfloat16) + return torch.randn(16, 64, device="cuda", dtype=torch.bfloat16) + + +def _identity_module_forward(module_name, module, x): + if module_name == "GroupedLinear": + m_splits = torch.tensor([8, 8], device="cuda", dtype=torch.int32) + return module(x, m_splits=m_splits) + return module(x) + + +def _fwd_bwd_module(module_name, model, x, recipe=None): + x = x.clone().detach().requires_grad_(True) + if recipe is not None: + with te.autocast(enabled=True, recipe=recipe): + y = _identity_module_forward(module_name, model, x) + else: + y = _identity_module_forward(module_name, model, x) + torch.manual_seed(99) + target = torch.randn_like(y) + loss = torch.nn.functional.mse_loss(y, target) + loss.backward() + wgrads = [p.grad.detach().clone() for p in model.parameters() if p.grad is not None] + return y.detach().clone(), x.grad.detach().clone(), wgrads + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +class TestIdentityTEModuleCoverage: + """All-Identity recipes should route every TE module through HP-compatible paths.""" + + @pytest.mark.parametrize("module_name", _IDENTITY_MODULE_NAMES) + @pytest.mark.parametrize( + "qfactory", + [ + pytest.param(identity_all_factory, id="plain_identity"), + pytest.param(hybrid_all_identity_factory, id="hybrid_identity"), + ], + ) + def test_identity_recipe_matches_bf16_bitwise(self, module_name, qfactory): + # Keep the GEMM topology identical. Identity-backed grad-output uses an + # unfused bgrad path, while the plain BF16 path may fuse bgrad with + # wgrad. Those mathematically equivalent kernels need not accumulate + # wgrad in the same order on every architecture. + ref, test = _make_identity_module_pair(module_name, seed=7300, bias=False) + x = _identity_module_input(module_name) + recipe = CustomRecipe(qfactory=qfactory) + + y_ref, dx_ref, wg_ref = _fwd_bwd_module(module_name, ref, x, recipe=None) + y_id, dx_id, wg_id = _fwd_bwd_module(module_name, test, x, recipe=recipe) + + # Identity is a high-precision passthrough. Entering the recipe context + # must not change any module's BF16 kernel selection or arithmetic. + torch.testing.assert_close(y_id, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_ref, rtol=0.0, atol=0.0) + assert len(wg_id) == len(wg_ref) + for g_id, g_ref in zip(wg_id, wg_ref): + torch.testing.assert_close(g_id, g_ref, rtol=0.0, atol=0.0) + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_grouped_linear_mxfp8_forward_identity_backward_matches_override(self): + def mxfp8_all_factory(role): # pylint: disable=unused-argument + return _mxfp8(tex.DType.kFloat8E4M3) + + def run(model, x, recipe): + x = x.detach().clone().requires_grad_(True) + m_splits = torch.tensor([32, 32], device="cuda", dtype=torch.int32) + with te.autocast(enabled=True, recipe=recipe): + y = model(x, m_splits=m_splits) + torch.manual_seed(9001) + target = torch.randn_like(y) + loss = torch.nn.functional.mse_loss(y, target) + loss.backward() + wgrads = [p.grad.detach().clone() for p in model.parameters() if p.grad is not None] + return y.detach().clone(), x.grad.detach().clone(), wgrads + + torch.manual_seed(8300) + ref = te.GroupedLinear(2, 64, 64, params_dtype=torch.bfloat16).cuda() + torch.manual_seed(8301) + test = te.GroupedLinear(2, 64, 64, params_dtype=torch.bfloat16).cuda() + with torch.no_grad(): + for p_test, p_ref in zip(test.parameters(), ref.parameters()): + p_test.copy_(p_ref) + + torch.manual_seed(8302) + x = torch.randn(64, 64, device="cuda", dtype=torch.bfloat16) + y_bo, dx_bo, wg_bo = run( + ref, + x, + CustomRecipe(qfactory=mxfp8_all_factory, backward_override="high_precision"), + ) + y_id, dx_id, wg_id = run( + test, + x, + CustomRecipe(qfactory=_hybrid_quantized_fwd_identity_bwd_factory("mxfp8")), + ) + + torch.testing.assert_close(y_id, y_bo, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_bo, rtol=0.0, atol=0.0) + assert len(wg_id) == len(wg_bo) + for g_id, g_bo in zip(wg_id, wg_bo): + torch.testing.assert_close(g_id, g_bo, rtol=0.0, atol=0.0) + + +class TestIdentityHybridFormatProtocols: + SHAPE = (256, 256) + OFFLOAD_SHAPE = (1024, 1024) + + @pytest.mark.parametrize("format_name", _HYBRID_IDENTITY_FORMATS) + def test_save_restore_keeps_identity_direction_exact(self, format_name): + torch.manual_seed(401) + x = torch.randn(*self.SHAPE, device="cuda", dtype=torch.bfloat16) + q = HybridQuantizer( + rowwise_quantizer=_format_quantizer(format_name), + columnwise_quantizer=IdentityQuantizer(), + ) + hybrid = q.quantize(x) + expected_row = hybrid._rowwise_storage.dequantize().clone() + expected_col = x.clone() + + tensors, obj = hybrid.prepare_for_saving() + leftover = obj.restore_from_saved(tensors) + + assert leftover == [] + assert isinstance(hybrid._columnwise_storage, IdentityTensorStorage) + torch.testing.assert_close( + hybrid._columnwise_storage.dequantize(), expected_col, rtol=0.0, atol=0.0 + ) + torch.testing.assert_close( + hybrid._rowwise_storage.dequantize(), expected_row, rtol=0.0, atol=0.0 + ) + + @pytest.mark.parametrize("format_name", _HYBRID_IDENTITY_FORMATS) + def test_cpu_offload_keeps_identity_direction_exact(self, format_name): + torch.manual_seed(402) + x = torch.randn(*self.OFFLOAD_SHAPE, device="cuda", dtype=torch.bfloat16) + q = HybridQuantizer( + rowwise_quantizer=_format_quantizer(format_name), + columnwise_quantizer=IdentityQuantizer(), + ) + hybrid = q.quantize(x) + expected_row = hybrid._rowwise_storage.dequantize().clone() + + reloaded = _offload_roundtrip(hybrid) + + assert isinstance(reloaded, HybridQuantizedTensor) + assert isinstance(reloaded._columnwise_storage, IdentityTensorStorage) + torch.testing.assert_close(reloaded._columnwise_storage.dequantize(), x, rtol=0.0, atol=0.0) + torch.testing.assert_close( + reloaded._rowwise_storage.dequantize(), expected_row, rtol=0.0, atol=0.0 + ) + + @pytest.mark.parametrize("format_name", ["mxfp8", "float8_blockwise", "nvfp4"]) + def test_quantize_master_weights_per_block_hybrid_identity_rejected(self, format_name): + if format_name == "mxfp8" and not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") + if format_name == "float8_blockwise" and not fp8_block_scaling_available: + pytest.skip(f"Float8Blockwise: {reason_for_no_fp8_block_scaling}") + if format_name == "nvfp4" and not (fp8_available and nvfp4_available): + pytest.skip(f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}") + + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + x = torch.randn(*self.SHAPE, device="cuda", dtype=torch.bfloat16) + q = HybridQuantizer( + rowwise_quantizer=_format_quantizer(format_name), + columnwise_quantizer=IdentityQuantizer(), + ) + weight = q.quantize(x) + master = torch.randn_like(x, dtype=torch.float32).reshape(-1).contiguous() + + with pytest.raises(NotImplementedError, match="HybridQuantizer"): + quantize_master_weights([weight], [master], [0], group=group) + + @pytest.mark.parametrize("format_name", _HYBRID_IDENTITY_RECOMPUTE_FORMATS) + @pytest.mark.parametrize("use_reentrant", [True, False]) + def test_activation_recompute_matches_no_checkpoint(self, format_name, use_reentrant): + recipe = CustomRecipe(qfactory=_hybrid_quantized_fwd_identity_bwd_factory(format_name)) + ref, test = _make_linears(128, 128, seed=440) + torch.manual_seed(441) + x = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16) + + y_ref, dx_ref, wg_ref = _fwd_bwd_checkpoint(ref, x, recipe, use_reentrant=None) + y_test, dx_test, wg_test = _fwd_bwd_checkpoint(test, x, recipe, use_reentrant=use_reentrant) + + torch.testing.assert_close(y_test, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_test, dx_ref, rtol=0.0, atol=0.0) + assert len(wg_test) == len(wg_ref) + for g_test, g_ref in zip(wg_test, wg_ref): + torch.testing.assert_close(g_test, g_ref, rtol=0.0, atol=0.0) + + +_ZOO_DEQUANTIZED_MODULES = [ + pytest.param("Linear", id="linear"), + pytest.param("LayerNormLinear", id="layernorm_linear"), + pytest.param("GroupedLinear", id="grouped_linear"), + pytest.param( + "LayerNormMLP", + marks=pytest.mark.xfail( + reason="LayerNormMLP does not support built-in backward_override=dequantized", + raises=AssertionError, + strict=True, + ), + id="layernorm_mlp", + ), + pytest.param("LayerNormLinearLinear", id="layernorm_linear_linear"), +] + +_ZOO_DEQUANTIZED_CASES = [ + pytest.param( + "mxfp8", + marks=pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}"), + id="mxfp8", + ), + pytest.param( + "nvfp4_row_scaled", + marks=pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}"), + id="nvfp4_row_scaled", + ), +] + + +def _make_zoo_dequantized_module(module_name, *, save_original_input=False): + hidden_size = 128 + if module_name == "Linear": + return te.Linear( + hidden_size, + hidden_size, + params_dtype=torch.bfloat16, + save_original_input=save_original_input, + ).cuda() + if module_name == "LayerNormLinear": + return te.LayerNormLinear(hidden_size, hidden_size, params_dtype=torch.bfloat16).cuda() + if module_name == "GroupedLinear": + return te.GroupedLinear( + 2, + hidden_size, + hidden_size, + params_dtype=torch.bfloat16, + save_original_input=save_original_input, + ).cuda() + if module_name == "LayerNormMLP": + return te.LayerNormMLP( + hidden_size, + 2 * hidden_size, + params_dtype=torch.bfloat16, + ).cuda() + if module_name == "LayerNormLinearLinear": + return torch.nn.Sequential( + te.LayerNormLinear(hidden_size, hidden_size, params_dtype=torch.bfloat16), + te.Linear( + hidden_size, + hidden_size, + params_dtype=torch.bfloat16, + save_original_input=save_original_input, + ), + ).cuda() + raise ValueError(module_name) + + +def _make_zoo_dequantized_module_pair(module_name): + torch.manual_seed(9400) + torch.cuda.manual_seed(9400) + ref = _make_zoo_dequantized_module(module_name, save_original_input=False) + torch.manual_seed(9401) + torch.cuda.manual_seed(9401) + test = _make_zoo_dequantized_module(module_name, save_original_input=True) + test.load_state_dict(ref.state_dict()) + return ref, test + + +def _zoo_dequantized_input(module_name): + torch.manual_seed(9402) + batch = 128 if module_name == "GroupedLinear" else 64 + return torch.randn(batch, 128, device="cuda", dtype=torch.bfloat16) + + +def _zoo_dequantized_forward(module_name, module, inp): + if module_name == "GroupedLinear": + m_splits = torch.tensor([64, 64], device="cuda", dtype=torch.int32) + return module(inp, m_splits=m_splits) + return module(inp) + + +def _fwd_bwd_zoo_dequantized_module(module_name, module, inp, recipe): + inp = inp.detach().clone().requires_grad_(True) + with te.autocast(enabled=True, recipe=recipe): + out = _zoo_dequantized_forward(module_name, module, inp) + torch.manual_seed(9403) + grad = torch.randn_like(out) + out.backward(grad) + param_grads = { + name: param.grad.detach().clone() + for name, param in module.named_parameters() + if param.grad is not None + } + return out.detach().clone(), inp.grad.detach().clone(), param_grads + + +def _mxfp8_all_qfactory(role): # pylint: disable=unused-argument + return _mxfp8(tex.DType.kFloat8E4M3) + + +def _zoo_dequantized_qfactory(case_name): + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + mxfp8_fwd_high_precision_bwd_factory, + nvfp4_row_scaled_fwd_high_precision_bwd_factory, + ) + + if case_name == "mxfp8": + return mxfp8_fwd_high_precision_bwd_factory + if case_name == "nvfp4_row_scaled": + return nvfp4_row_scaled_fwd_high_precision_bwd_factory + raise ValueError(case_name) + + +def _zoo_dequantized_recipes(case_name): + from transformer_engine.common import recipe as te_recipe + + if case_name == "mxfp8": + return ( + CustomRecipe(qfactory=_mxfp8_all_qfactory, backward_override="dequantized"), + CustomRecipe(qfactory=_zoo_dequantized_qfactory(case_name)), + ) + if case_name == "nvfp4_row_scaled": + ref_recipe = te_recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + disable_2d_quantization=True, + row_scaled_activation=True, + backward_override="dequantized", + ) + ref_recipe.fp4_quant_fwd_inp = te_recipe.QParams() + ref_recipe.fp4_quant_fwd_weight = te_recipe.QParams() + ref_recipe.fp4_quant_bwd_grad = te_recipe.QParams() + return ( + ref_recipe, + CustomRecipe(qfactory=_zoo_dequantized_qfactory(case_name)), + ) + raise ValueError(case_name) + + +def _assert_zoo_layernorm_mlp_role_semantics(case_name): + module = _make_zoo_dequantized_module("LayerNormMLP") + qfactory = _zoo_dequantized_qfactory(case_name) + + fwd_roles = module.get_quantizer_roles(fwd=True, num_quantizers=6) + fwd_gemm_roles = [ + role for role in fwd_roles if role is not None and role.tensor_type in ("input", "weight") + ] + assert len(fwd_gemm_roles) == 5 + for role in fwd_gemm_roles: + quantizer = qfactory(role) + assert isinstance(quantizer, HybridQuantizer) + assert quantizer.columnwise_source == "rowwise_dequantized" + assert isinstance(quantizer.columnwise_quantizer, IdentityQuantizer) + if case_name == "mxfp8": + assert isinstance(quantizer.rowwise_quantizer, MXFP8Quantizer) + else: + assert isinstance(quantizer.rowwise_quantizer, NVFP4Quantizer) + assert quantizer.rowwise_quantizer.row_scaled_nvfp4 == (role.tensor_type == "input") + + bwd_roles = module.get_quantizer_roles(fwd=False, num_quantizers=4) + grad_output_roles = [ + role for role in bwd_roles if role is not None and role.tensor_type == "grad_output" + ] + assert len(grad_output_roles) == 3 + for role in grad_output_roles: + assert isinstance(qfactory(role), IdentityQuantizer) + + +class TestZooHighPrecisionBackwardFactoryModuleCoverage: + """Zoo high-precision-backward factories should match base recipes across TE modules.""" + + @pytest.mark.parametrize("case_name", _ZOO_DEQUANTIZED_CASES) + @pytest.mark.parametrize("module_name", _ZOO_DEQUANTIZED_MODULES) + def test_matches_base_recipe_bitwise(self, case_name, module_name): + ref, test = _make_zoo_dequantized_module_pair(module_name) + inp = _zoo_dequantized_input(module_name) + ref_recipe, qfactory_recipe = _zoo_dequantized_recipes(case_name) + + y_ref, dx_ref, grads_ref = _fwd_bwd_zoo_dequantized_module( + module_name, ref, inp, ref_recipe + ) + y_test, dx_test, grads_test = _fwd_bwd_zoo_dequantized_module( + module_name, test, inp, qfactory_recipe + ) + + torch.testing.assert_close(y_test, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_test, dx_ref, rtol=0.0, atol=0.0) + assert grads_test.keys() == grads_ref.keys() + for name, grad_ref in grads_ref.items(): + torch.testing.assert_close( + grads_test[name], + grad_ref, + rtol=0.0, + atol=0.0, + msg=f"{case_name}/{module_name} grad mismatch for {name}", + ) + + @pytest.mark.parametrize("case_name", _ZOO_DEQUANTIZED_CASES) + def test_layernorm_mlp_runs_and_uses_dequantized_backward_roles(self, case_name): + _assert_zoo_layernorm_mlp_role_semantics(case_name) + + module = _make_zoo_dequantized_module("LayerNormMLP") + inp = _zoo_dequantized_input("LayerNormMLP") + recipe = CustomRecipe(qfactory=_zoo_dequantized_qfactory(case_name)) + out, dx, grads = _fwd_bwd_zoo_dequantized_module("LayerNormMLP", module, inp, recipe) + + assert torch.isfinite(out).all() + assert torch.isfinite(dx).all() + expected_grads = {name for name, param in module.named_parameters() if param.requires_grad} + assert grads.keys() == expected_grads + for grad in grads.values(): + assert torch.isfinite(grad).all() + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +class TestIdentityLinear: + """End-to-end te.Linear with Identity-based recipes.""" + + IN_F = 128 + OUT_F = 128 + BATCH = 64 + + def _input(self): + torch.manual_seed(7) + return torch.randn(self.BATCH, self.IN_F, device="cuda", dtype=torch.bfloat16) + + @pytest.mark.parametrize( + "qfactory", + [ + pytest.param(identity_all_factory, id="plain_identity"), + pytest.param(hybrid_all_identity_factory, id="hybrid_identity"), + ], + ) + def test_identity_bias_path_matches_bf16_bitwise(self, qfactory): + """The unfused Identity bgrad path must preserve exact bias semantics. + + Weight-gradient bitwise parity is covered with ``bias=False`` so both + sides use the same GEMM topology. This bias-enabled test separately + locks output, input-gradient, and bias-gradient parity to zero tolerance. + """ + ref, test = _make_linears(self.IN_F, self.OUT_F, bias=True) + x = self._input() + + y_ref, dx_ref, _ = _fwd_bwd(ref, x, recipe=None) + y_id, dx_id, _ = _fwd_bwd(test, x, recipe=CustomRecipe(qfactory=qfactory)) + + torch.testing.assert_close(y_id, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(test.bias.grad, ref.bias.grad, rtol=0.0, atol=0.0) + + def test_whole_layer_hp_matches_bf16_bitwise(self): + """Identity for every slot => all GEMMs high precision => bitwise-equal + to a plain BF16 te.Linear (no autocast).""" + ref, test = _make_linears(self.IN_F, self.OUT_F, bias=False) + x = self._input() + + y_ref, dx_ref, wg_ref = _fwd_bwd(ref, x, recipe=None) + y_id, dx_id, wg_id = _fwd_bwd(test, x, recipe=CustomRecipe(qfactory=identity_all_factory)) + + torch.testing.assert_close(y_id, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_ref, rtol=0.0, atol=0.0) + assert len(wg_id) == len(wg_ref) + for g_id, g_ref in zip(wg_id, wg_ref): + torch.testing.assert_close(g_id, g_ref, rtol=0.0, atol=0.0) + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_fwd_hp_bwd_fp8_forward_bitwise(self): + """High-precision forward must be bitwise-equal to BF16 forward; the + backward runs in FP8 (finite, close to BF16 within a loose tolerance).""" + ref, test = _make_linears(self.IN_F, self.OUT_F) + x = self._input() + + y_ref, dx_ref, wg_ref = _fwd_bwd(ref, x, recipe=None) + y_h, dx_h, wg_h = _fwd_bwd(test, x, recipe=CustomRecipe(qfactory=fwd_hp_bwd_fp8_factory)) + + # Forward is high precision -> bitwise equal. + torch.testing.assert_close(y_h, y_ref, rtol=0.0, atol=0.0) + # Backward is FP8 (E4M3 weight col, E5M2 grad) -> relative L2 error vs the + # BF16 reference reflects pure FP8 quant noise. Measured: dgrad ~5.7e-2, + # weight-grad ~5.8e-2 (E4M3 ~6% step). Bound 7e-2 keeps a small margin. + assert torch.isfinite(dx_h).all() + assert _rel_l2_error(dx_h, dx_ref) < 7e-2 + for g, g_ref in zip(wg_h, wg_ref): + assert torch.isfinite(g).all() + if g.dim() == 1: + # Bias grad = sum(dY) is computed in high precision (dY is bitwise + # identical since the forward is bitwise), so it must match exactly. + torch.testing.assert_close(g, g_ref, rtol=0.0, atol=0.0) + else: + assert _rel_l2_error(g, g_ref) < 7e-2 + + def test_fwd_fp8_bwd_hp_runs_and_backward_high_precision(self): + """FP8 forward + high-precision backward. Forward differs from BF16 + (quantized), backward GEMMs run in high precision.""" + ref, test = _make_linears(self.IN_F, self.OUT_F) + x = self._input() + + y_ref, dx_ref, _ = _fwd_bwd(ref, x, recipe=None) + y_q, dx_q, wg_q = _fwd_bwd(test, x, recipe=CustomRecipe(qfactory=fwd_fp8_bwd_hp_factory)) + + # Forward is FP8 (E4M3) -> relative L2 error vs BF16 is the quant noise. + # Measured ~3.7e-2; bound 5e-2. + assert torch.isfinite(y_q).all() + assert _rel_l2_error(y_q, y_ref) < 5e-2 + # Backward GEMMs run in high precision. dgrad differs from the BF16 + # reference only because the FP8 forward perturbs dY; measured ~1.2e-2, + # bound 3e-2 (the bitwise HP-backward guarantee is locked by the + # backward_override equivalence test below). + assert torch.isfinite(dx_q).all() + assert _rel_l2_error(dx_q, dx_ref) < 3e-2 + for g in wg_q: + assert torch.isfinite(g).all() + + def test_hybrid_all_identity_matches_bf16_bitwise(self): + """All-Identity through the *hybrid* container must be bitwise-equal to a + plain BF16 te.Linear. Complements the non-hybrid whole-layer-HP test: this + exercises HybridQuantizedTensor with Identity sub-storages in both + directions and the per-operand unwrap of every GEMM.""" + ref, test = _make_linears(self.IN_F, self.OUT_F, bias=False) + x = self._input() + + y_ref, dx_ref, wg_ref = _fwd_bwd(ref, x, recipe=None) + y_id, dx_id, wg_id = _fwd_bwd( + test, x, recipe=CustomRecipe(qfactory=hybrid_all_identity_factory) + ) + + torch.testing.assert_close(y_id, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_ref, rtol=0.0, atol=0.0) + assert len(wg_id) == len(wg_ref) + for g_id, g_ref in zip(wg_id, wg_ref): + torch.testing.assert_close(g_id, g_ref, rtol=0.0, atol=0.0) + + def test_identity_reproduces_backward_override_high_precision_bitwise(self): + """The per-direction Identity machinery must reproduce + ``backward_override="high_precision"`` **bitwise**. + + Both runs quantize the forward to the same FP8 (current scaling) and run + the backward in high precision against the original operands. The Identity + path expresses this per-tensor (weight/input = Hybrid(row=FP8, col=Identity), + grad = Identity); the reference uses the global ``backward_override`` knob. + Identical forward FP8 + identical original HP backward operands => bitwise. + """ + ref, test = _make_linears(self.IN_F, self.OUT_F, bias=False) + x = self._input() + + y_bo, dx_bo, wg_bo = _fwd_bwd( + ref, + x, + recipe=CustomRecipe(qfactory=fp8_fwd_factory, backward_override="high_precision"), + ) + y_id, dx_id, wg_id = _fwd_bwd(test, x, recipe=CustomRecipe(qfactory=fwd_fp8_bwd_hp_factory)) + + torch.testing.assert_close(y_id, y_bo, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_bo, rtol=0.0, atol=0.0) + assert len(wg_id) == len(wg_bo) + for g_id, g_bo in zip(wg_id, wg_bo): + torch.testing.assert_close(g_id, g_bo, rtol=0.0, atol=0.0) + + def test_identity_reproduces_backward_override_dequantized_bitwise(self): + """Per-direction Identity must reproduce ``backward_override=dequantized``. + + Both runs quantize forward to the same FP8 values. The reference saves + rowwise fprop payloads and dequantizes them for high-precision backward; + the hybrid path stores Identity columnwise data sourced from the rowwise + dequantized value and uses Identity grad tensors. + """ + ref, test = _make_linears(self.IN_F, self.OUT_F, bias=False) + x = self._input() + + y_bo, dx_bo, wg_bo = _fwd_bwd( + ref, + x, + recipe=CustomRecipe(qfactory=fp8_fwd_factory, backward_override="dequantized"), + ) + y_id, dx_id, wg_id = _fwd_bwd( + test, x, recipe=CustomRecipe(qfactory=fwd_fp8_bwd_rowwise_dequantized_hp_factory) + ) + + torch.testing.assert_close(y_id, y_bo, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_bo, rtol=0.0, atol=0.0) + assert len(wg_id) == len(wg_bo) + for g_id, g_bo in zip(wg_id, wg_bo): + torch.testing.assert_close(g_id, g_bo, rtol=0.0, atol=0.0) + + def test_identity_matches_bf16_multistep_training_bitwise(self): + """Multi-step SGD: an all-Identity recipe must track a plain BF16 + te.Linear bitwise across optimizer steps (no drift from workspace caching + or any hidden state).""" + ref, test = _make_linears(self.IN_F, self.OUT_F, bias=False) + opt_ref = torch.optim.SGD(ref.parameters(), lr=0.1) + opt_test = torch.optim.SGD(test.parameters(), lr=0.1) + recipe = CustomRecipe(qfactory=identity_all_factory) + + for step in range(4): + torch.manual_seed(1000 + step) + x = torch.randn(self.BATCH, self.IN_F, device="cuda", dtype=torch.bfloat16) + torch.manual_seed(2000 + step) + target = torch.randn(self.BATCH, self.OUT_F, device="cuda", dtype=torch.bfloat16) + + opt_ref.zero_grad() + y_ref = ref(x) + loss_ref = torch.nn.functional.mse_loss(y_ref, target) + loss_ref.backward() + opt_ref.step() + + opt_test.zero_grad() + with te.autocast(enabled=True, recipe=recipe): + y_test = test(x) + loss_test = torch.nn.functional.mse_loss(y_test, target) + loss_test.backward() + opt_test.step() + + torch.testing.assert_close(y_test, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(loss_test, loss_ref, rtol=0.0, atol=0.0) + for p_test, p_ref in zip(test.parameters(), ref.parameters()): + torch.testing.assert_close(p_test, p_ref, rtol=0.0, atol=0.0, msg=f"step {step}") + + def test_quantized_model_init_identity_matches_bf16_bitwise(self): + """Persistent Identity params from quantized_model_init should match BF16 exactly.""" + torch.manual_seed(314) + ref = te.Linear(self.IN_F, self.OUT_F, bias=False, params_dtype=torch.bfloat16).cuda() + recipe = CustomRecipe(qfactory=identity_all_factory) + torch.manual_seed(2718) + with te.quantized_model_init(enabled=True, recipe=recipe): + test = te.Linear(self.IN_F, self.OUT_F, bias=False, params_dtype=torch.bfloat16).cuda() + with torch.no_grad(): + for p_test, p_ref in zip(test.parameters(), ref.parameters()): + assert isinstance(p_test, IdentityTensor) + p_test.copy_(p_ref) + + x = self._input() + y_ref, dx_ref, wg_ref = _fwd_bwd(ref, x, recipe=None) + y_id, dx_id, wg_id = _fwd_bwd(test, x, recipe=recipe) + + torch.testing.assert_close(y_id, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_ref, rtol=0.0, atol=0.0) + for g_id, g_ref in zip(wg_id, wg_ref): + torch.testing.assert_close(g_id, g_ref, rtol=0.0, atol=0.0) + + def test_quantized_model_init_identity_training_loss_decreases_bitwise(self): + """All-Identity quantized params train like BF16 and loss decreases.""" + torch.manual_seed(777) + ref = te.Linear(self.IN_F, self.OUT_F, bias=False, params_dtype=torch.bfloat16).cuda() + recipe = CustomRecipe(qfactory=identity_all_factory) + torch.manual_seed(888) + with te.quantized_model_init(enabled=True, recipe=recipe): + test = te.Linear(self.IN_F, self.OUT_F, bias=False, params_dtype=torch.bfloat16).cuda() + with torch.no_grad(): + for p_test, p_ref in zip(test.parameters(), ref.parameters()): + assert isinstance(p_test, IdentityTensor) + p_test.copy_(p_ref) + + torch.manual_seed(909) + x = torch.randn(self.BATCH, self.IN_F, device="cuda", dtype=torch.bfloat16) + target = torch.zeros(self.BATCH, self.OUT_F, device="cuda", dtype=torch.bfloat16) + opt_ref = torch.optim.SGD(ref.parameters(), lr=0.1) + opt_test = torch.optim.SGD(test.parameters(), lr=0.1) + losses_ref = [] + losses_test = [] + + for _ in range(5): + opt_ref.zero_grad() + y_ref = ref(x) + loss_ref = torch.nn.functional.mse_loss(y_ref, target) + loss_ref.backward() + opt_ref.step() + losses_ref.append(loss_ref.detach().clone()) + + opt_test.zero_grad() + with te.autocast(enabled=True, recipe=recipe): + y_test = test(x) + loss_test = torch.nn.functional.mse_loss(y_test, target) + loss_test.backward() + opt_test.step() + losses_test.append(loss_test.detach().clone()) + + torch.testing.assert_close(y_test, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(loss_test, loss_ref, rtol=0.0, atol=0.0) + for p_test, p_ref in zip(test.parameters(), ref.parameters()): + torch.testing.assert_close(p_test.dequantize(), p_ref, rtol=0.0, atol=0.0) + + assert all( + losses_ref[i + 1].item() < losses_ref[i].item() for i in range(len(losses_ref) - 1) + ), f"BF16 loss did not strictly decrease: {[x.item() for x in losses_ref]}" + for loss_test, loss_ref in zip(losses_test, losses_ref): + torch.testing.assert_close(loss_test, loss_ref, rtol=0.0, atol=0.0) + + @pytest.mark.parametrize("use_reentrant", [True, False]) + def test_identity_activation_recompute_matches_bf16_bitwise(self, use_reentrant): + """All-Identity recompute should be exactly the BF16 no-checkpoint path.""" + ref, test = _make_linears(self.IN_F, self.OUT_F, seed=4242, bias=False) + recipe = CustomRecipe(qfactory=identity_all_factory) + x = self._input() + + y_ref, dx_ref, wg_ref = _fwd_bwd(ref, x, recipe=None) + y_id, dx_id, wg_id = _fwd_bwd_checkpoint(test, x, recipe, use_reentrant=use_reentrant) + + torch.testing.assert_close(y_id, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_ref, rtol=0.0, atol=0.0) + for g_id, g_ref in zip(wg_id, wg_ref): + torch.testing.assert_close(g_id, g_ref, rtol=0.0, atol=0.0) + + def test_quantized_model_init_identity_state_dict_save_load_exact(self): + recipe = CustomRecipe(qfactory=identity_all_factory) + torch.manual_seed(5151) + with te.quantized_model_init(enabled=True, recipe=recipe): + model = te.Linear(64, 64, bias=False, params_dtype=torch.bfloat16).cuda() + + torch.manual_seed(5152) + x = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(), te.autocast(enabled=True, recipe=recipe): + out_before = model(x) + + buffer = io.BytesIO() + torch.save(model.state_dict(), buffer) + buffer.seek(0) + + with te.quantized_model_init(enabled=True, recipe=recipe): + model2 = te.Linear(64, 64, bias=False, params_dtype=torch.bfloat16).cuda() + model2.load_state_dict(torch.load(buffer, weights_only=True)) + + with torch.no_grad(), te.autocast(enabled=True, recipe=recipe): + out_after = model2(x) + + assert isinstance(model2.weight, IdentityTensor) + torch.testing.assert_close(out_after, out_before, rtol=0.0, atol=0.0) + + def test_load_bf16_state_dict_into_identity_model_exact(self): + recipe = CustomRecipe(qfactory=identity_all_factory) + torch.manual_seed(6161) + ref = te.Linear(64, 64, bias=False, params_dtype=torch.bfloat16).cuda() + with te.quantized_model_init(enabled=True, recipe=recipe): + model = te.Linear(64, 64, bias=False, params_dtype=torch.bfloat16).cuda() + + model.load_state_dict(ref.state_dict()) + assert isinstance(model.weight, IdentityTensor) + torch.testing.assert_close(model.weight.dequantize(), ref.weight, rtol=0.0, atol=0.0) + + torch.manual_seed(6162) + x = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + out_ref = ref(x) + with te.autocast(enabled=True, recipe=recipe): + out_id = model(x) + + torch.testing.assert_close(out_id, out_ref, rtol=0.0, atol=0.0) diff --git a/tests/pytorch/test_mhc.py b/tests/pytorch/test_mhc.py index 541ce9a8c2..6e1f556678 100644 --- a/tests/pytorch/test_mhc.py +++ b/tests/pytorch/test_mhc.py @@ -9,18 +9,20 @@ from utils import reset_rng_states from transformer_engine.pytorch.triton.mhc import ( + ENFORCE_DETERMINISTIC, mhc_fused_sinkhorn, mhc_fused_scale, mhc_fused_aggregate, mhc_fused_expand_combine, mhc_fused_projection, + mhc_generate_mix_and_aggregate, ) # Disable TF32 for matmul to ensure consistency between the fused and reference implementations torch.backends.cuda.matmul.allow_tf32 = False -def mhc_projection_ref(x, phi): +def mhc_projection_ref(x, phi, norm_weight): """ Reference operator for mHC's projection building operation. @@ -29,19 +31,20 @@ def mhc_projection_ref(x, phi): - phi_pre: (n, nC) - phi_post: (n, nC) - phi_res: (n^2, nC) + norm_weight: (nC,) or None, if not None, apply element-wise multiplication to phi before projection n: number of Hyper Connection streams C: hidden dimension per stream """ - x_dtype = x.dtype - x = x.to(torch.float32) - phi = phi.to(torch.float32) - Hs = x @ phi.T # (M, 2n + n^2) + x_fp64 = x.to(torch.float64) + ms = (x_fp64 * x_fp64).mean(dim=1) - x_fp32 = x.to(torch.float32) # Use fp32 for better numerical stability in variance calculation - ms = (x_fp32 * x_fp32).mean(dim=1) + phi_fp64 = phi.to(torch.float64) + if norm_weight is not None: + phi_fp64 = phi_fp64 * norm_weight.to(torch.float64)[None, :] + Hs = x_fp64 @ phi_fp64.T # (M, 2n + n^2) - return Hs.to(x_dtype), ms + return Hs, ms def mhc_scale_ref(H, alpha, beta, ms, n): @@ -139,9 +142,9 @@ def mhc_aggregate_ref(x, H_pre, n): s, b, C, n = x.shape H_pre = H_pre.view(s, b, n, 1) - out = (x @ H_pre).view(s, b, C) + out = (x.to(H_pre.dtype) @ H_pre).view(s, b, C) - return out + return out.to(x.dtype) def mhc_expand_combine_ref(f, bias, H_post, x, H_res, n): @@ -267,38 +270,63 @@ def get_tols(dtype): @pytest.mark.parametrize("cfg", mhc_configs, ids=MHCConfig.desc) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"]) -def test_mhc_projection(cfg: MHCConfig, dtype): +@pytest.mark.parametrize( + "dtypes", + [ + (torch.float32, torch.float32), + (torch.bfloat16, torch.bfloat16), + (torch.bfloat16, torch.float32), + ], + ids=["x_fp32_phi_fp32", "x_bf16_phi_bf16", "x_bf16_phi_fp32"], +) +@pytest.mark.parametrize("has_norm_weight", [False, True], ids=["no_norm_weight", "norm_weight"]) +@pytest.mark.parametrize("use_split_k", [True, False], ids=["split_k", "no_split_k"]) +def test_mhc_projection(cfg: MHCConfig, dtypes, has_norm_weight, use_split_k): reset_rng_states() + if ENFORCE_DETERMINISTIC and use_split_k: + pytest.skip("Split-K is not deterministic, skip the test under deterministic mode") + s, b, C, n = cfg.s, cfg.b, cfg.C, cfg.n nC = n * C N = 2 * n + n * n - tols = get_tols(dtype) + x_dtype = dtypes[0] + phi_dtype = dtypes[1] + tols = get_tols(x_dtype) use_tf32 = False - x = torch.randn(s * b, nC, device="cuda", requires_grad=True, dtype=dtype) - phi = torch.randn(N, nC, dtype=dtype, requires_grad=True, device="cuda") - + x = torch.randn(s * b, nC, device="cuda", requires_grad=True, dtype=x_dtype) + phi = torch.randn(N, nC, dtype=phi_dtype, requires_grad=True, device="cuda") x_ref = x.detach().clone().requires_grad_(True) phi_ref = phi.detach().clone().requires_grad_(True) - ref_out_Hs, ref_out_ms = mhc_projection_ref(x_ref, phi_ref) - fused_out_Hs_padded, fused_out_ms = mhc_fused_projection(x, phi, use_tf32) + if has_norm_weight: + norm_weight = torch.randn(nC, device="cuda", requires_grad=True, dtype=x_dtype) + norm_weight_ref = norm_weight.detach().clone().requires_grad_(True) + else: + norm_weight = None + norm_weight_ref = None + + ref_out_Hs, ref_out_ms = mhc_projection_ref(x_ref, phi_ref, norm_weight_ref) + fused_out_Hs_padded, fused_out_ms = mhc_fused_projection( + x, phi, norm_weight=norm_weight, use_tf32=use_tf32, use_split_k=use_split_k + ) fused_out_Hs = fused_out_Hs_padded[:, :N] - torch.testing.assert_close(fused_out_Hs, ref_out_Hs, **tols) - torch.testing.assert_close(fused_out_ms, ref_out_ms, **tols) + torch.testing.assert_close(fused_out_Hs.double(), ref_out_Hs, **tols) + torch.testing.assert_close(fused_out_ms.double(), ref_out_ms, **tols) (ref_out_Hs.sum() + ref_out_ms.sum()).backward() (fused_out_Hs.sum() + fused_out_ms.sum()).backward() torch.testing.assert_close(x.grad, x_ref.grad, **tols) torch.testing.assert_close(phi.grad, phi_ref.grad, **tols) + if has_norm_weight: + torch.testing.assert_close(norm_weight.grad, norm_weight_ref.grad, **tols) @pytest.mark.parametrize("cfg", mhc_configs, ids=MHCConfig.desc) -@pytest.mark.parametrize("dtype", [torch.float32], ids=["fp32"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"]) def test_mhc_scale(cfg: MHCConfig, dtype): reset_rng_states() @@ -329,56 +357,79 @@ def test_mhc_scale(cfg: MHCConfig, dtype): torch.cat([fused_out[i] for i in range(3)], dim=-1).sum().backward() torch.testing.assert_close(H_padded.grad[:, :N], H_ref.grad, **tols) + torch.testing.assert_close(ms.grad, ms_ref.grad, **tols) torch.testing.assert_close(alpha.grad, alpha_ref.grad, **tols) torch.testing.assert_close(beta.grad, beta_ref.grad, **tols) - torch.testing.assert_close(ms.grad, ms_ref.grad, **tols) @pytest.mark.parametrize("cfg", mhc_configs, ids=MHCConfig.desc) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"]) -def test_mhc_combined(cfg: MHCConfig, dtype): +@pytest.mark.parametrize( + "dtypes", + [ + (torch.float32, torch.float32), + (torch.bfloat16, torch.bfloat16), + (torch.bfloat16, torch.float32), + ], + ids=["x_fp32_phi_fp32", "x_bf16_phi_bf16", "x_bf16_phi_fp32"], +) +@pytest.mark.parametrize("has_norm_weight", [False, True], ids=["no_norm_weight", "norm_weight"]) +@pytest.mark.parametrize("use_split_k", [True, False], ids=["split_k", "no_split_k"]) +def test_mhc_rmsnorm(cfg: MHCConfig, dtypes, has_norm_weight, use_split_k): + # Validate the fused (split-order) kernel against RMSNorm applied in the correct order, + # computed in fp64 via F.rms_norm as the oracle. In exact arithmetic the kernel's + # "matmul then divide by rms" equals "rms_norm(x) then matmul", so fp64 correct-order is + # ground truth for the fp32 kernel and is independent of the kernel's own reimplementation. reset_rng_states() + if ENFORCE_DETERMINISTIC and use_split_k: + pytest.skip("Split-K is not deterministic, skip the test under deterministic mode") + s, b, C, n = cfg.s, cfg.b, cfg.C, cfg.n N = 2 * n + n * n nC = n * C - tols = get_tols(dtype) + x_dtype = dtypes[0] + phi_dtype = dtypes[1] + tols = get_tols(x_dtype) use_tf32 = False - x = torch.randn(s * b, nC, device="cuda", requires_grad=True, dtype=dtype) - phi = torch.randn(N, nC, dtype=dtype, requires_grad=True, device="cuda") - - alpha = torch.randn(3, device="cuda", requires_grad=True, dtype=dtype) - beta = torch.randn(1, 2 * n + n * n, device="cuda", requires_grad=True, dtype=dtype) + x = torch.randn(s * b, nC, device="cuda", requires_grad=True, dtype=x_dtype) + phi = torch.randn(N, nC, dtype=phi_dtype, requires_grad=True, device="cuda") + alpha = torch.randn(3, device="cuda", requires_grad=True, dtype=phi_dtype) + beta = torch.randn(1, 2 * n + n * n, device="cuda", requires_grad=True, dtype=phi_dtype) x_ref = x.detach().clone().requires_grad_(True) phi_ref = phi.detach().clone().requires_grad_(True) - alpha_ref = alpha.detach().clone().requires_grad_(True) beta_ref = beta.detach().clone().requires_grad_(True) - ref_out_H, ref_out_r = mhc_projection_ref(x_ref, phi_ref) - fused_out_H_padded, fused_out_r = mhc_fused_projection(x, phi, use_tf32) + if has_norm_weight: + norm_weight = torch.randn(nC, device="cuda", requires_grad=True, dtype=x_dtype) + norm_weight_ref = norm_weight.detach().clone().requires_grad_(True) + else: + norm_weight = None + norm_weight_ref = None - ref_H_pre, ref_H_post, ref_H_res = mhc_scale_ref( - ref_out_H[:, :N], alpha_ref, beta_ref, ref_out_r, n + fused_out_H_padded, fused_out_r = mhc_fused_projection( + x, phi, norm_weight=norm_weight, use_tf32=use_tf32, use_split_k=use_split_k ) fused_H_pre, fused_H_post, fused_H_res = mhc_fused_scale( fused_out_H_padded, alpha, beta, fused_out_r, n ) - def mhc_combined(x_ref, phi_ref, alpha_ref, beta_ref): - dtype = x_ref.dtype - x_ref = x_ref.to(torch.float32) - phi_ref = phi_ref.to(torch.float32) - alpha_ref = alpha_ref.to(torch.float32) - beta_ref = beta_ref.to(torch.float32) - + def mhc_combined(x_ref, phi_ref, alpha_ref, beta_ref, norm_weight_ref): # Check if after spliting RMSNorm to two steps in projection and scaling, - # theresult is close to applying RMSNorm in the correct order - x_rmsnorm = F.rms_norm(x_ref, normalized_shape=(nC,)) - H = x_rmsnorm @ phi_ref.T + # the result is close to applying RMSNorm in the correct order. + # Run RMSNorm in fp32 so the bf16 case has the same precision pattern as the + # kernel/ref (F.rms_norm on bf16 input would round x_rmsnorm back to bf16). + eps = torch.finfo(torch.float64).eps + norm_weight_fp64 = ( + norm_weight_ref.to(torch.float64) if norm_weight_ref is not None else None + ) + x_rmsnorm = F.rms_norm( + x_ref.to(torch.float64), normalized_shape=(nC,), weight=norm_weight_fp64, eps=eps + ) + H = x_rmsnorm @ phi_ref.T.to(torch.float64) H_pre = H[:, :n] H_post = H[:, n : 2 * n] H_res = H[:, 2 * n :] @@ -391,19 +442,82 @@ def mhc_combined(x_ref, phi_ref, alpha_ref, beta_ref): out_post = 2 * out_post.sigmoid() out_res = out_res - return out_pre.to(dtype), out_post.to(dtype), out_res.to(dtype) + return out_pre, out_post, out_res # Return in FP32 to match the kernel's behavior combined_H_pre, combined_H_post, combined_H_res = mhc_combined( - x_ref, phi_ref, alpha_ref, beta_ref + x_ref, phi_ref, alpha_ref, beta_ref, norm_weight_ref + ) + + # We only need to compare F.rmsnorm and our fused implementation where we split RMSNorm into two steps + torch.testing.assert_close(combined_H_pre, fused_H_pre.to(torch.float64), **tols) + torch.testing.assert_close(combined_H_post, fused_H_post.to(torch.float64), **tols) + torch.testing.assert_close(combined_H_res, fused_H_res.to(torch.float64), **tols) + + +@pytest.mark.skipif( + not ENFORCE_DETERMINISTIC, + reason=( + "Skipped when NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 due to atomic_add nondeterminism" + " introducing too much error across multiple ops." + ), +) +@pytest.mark.parametrize("cfg", mhc_configs, ids=MHCConfig.desc) +@pytest.mark.parametrize("dtype", [torch.float32], ids=["fp32"]) +def test_mhc_fuse_grad_acc(cfg: MHCConfig, dtype): + # Skip bf16 tests since in the unfused path the we accumulate 3 bf16 gradients, whereas in the fused path + # we accumulate 3 fp32 gradients and then cast to bf16 in the end, which causes two paths to have different precision patterns + + reset_rng_states() + + s, b, C, n = cfg.s, cfg.b, cfg.C, cfg.n + N = 2 * n + n * n + nC = n * C + + # For non-deterministic tests, we use a looser tolerance since atomic add introduces greater error + tols = dict(atol=1e-4, rtol=1e-4) if ENFORCE_DETERMINISTIC else dict(atol=5e-2, rtol=5e-2) + use_tf32 = False + + x = torch.randn(s, b, C, n, device="cuda", requires_grad=True, dtype=dtype) + phi = torch.randn(N, nC, dtype=dtype, requires_grad=True, device="cuda") + + alpha = torch.randn(3, device="cuda", requires_grad=True, dtype=dtype) + beta = torch.randn(1, 2 * n + n * n, device="cuda", requires_grad=True, dtype=dtype) + x_ref = x.detach().clone().requires_grad_(True) + phi_ref = phi.detach().clone().requires_grad_(True) + + alpha_ref = alpha.detach().clone().requires_grad_(True) + beta_ref = beta.detach().clone().requires_grad_(True) + + def end_to_end(x, phi, alpha, beta, fused_grad_x_acc): + fused_grad_x_acc_buffer = None + if fused_grad_x_acc: + fused_grad_x_acc_buffer = torch.empty_like(x, dtype=torch.float32) + aggregated, H_post, H_res = mhc_generate_mix_and_aggregate( + x, phi, alpha, beta, None, use_tf32, fused_grad_x_acc_buffer + ) + expanded_combined = mhc_fused_expand_combine( + aggregated, + None, + H_post, + x, + H_res, + n, + False, + fused_grad_x_acc_buffer, + ) + + return expanded_combined + + expanded_combined_fuse_grad = end_to_end( + x_ref, phi_ref, alpha_ref, beta_ref, fused_grad_x_acc=True ) + expanded_combined_no_fuse_grad = end_to_end(x, phi, alpha, beta, fused_grad_x_acc=False) - torch.testing.assert_close(combined_H_pre, ref_H_pre, **tols) - torch.testing.assert_close(combined_H_post, ref_H_post, **tols) - torch.testing.assert_close(combined_H_res, ref_H_res, **tols) + grad_output = torch.randn_like(expanded_combined_fuse_grad) + expanded_combined_fuse_grad.backward(grad_output) + expanded_combined_no_fuse_grad.backward(grad_output) - torch.testing.assert_close(combined_H_pre, fused_H_pre, **tols) - torch.testing.assert_close(combined_H_post, fused_H_post, **tols) - torch.testing.assert_close(combined_H_res, fused_H_res, **tols) + torch.testing.assert_close(x.grad, x_ref.grad, **tols) @pytest.mark.parametrize("cfg", mhc_configs, ids=MHCConfig.desc) @@ -446,7 +560,7 @@ def test_mhc_aggregate(cfg: MHCConfig, dtype): H_pre_ref = H_pre.detach().clone().requires_grad_(True) ref_out = mhc_aggregate_ref(x_ref, H_pre_ref, n) - fused_out = mhc_fused_aggregate(x, H_pre, n, False) + fused_out = mhc_fused_aggregate(x, H_pre, n, use_tf32=False) torch.testing.assert_close(fused_out, ref_out, **tols) @@ -482,7 +596,7 @@ def test_mhc_expand_combine(cfg: MHCConfig, dtype, with_bias): H_res_ref = H_res.detach().clone().requires_grad_(True) ref_out = mhc_expand_combine_ref(f_ref, bias_ref, H_post_ref, x_ref, H_res_ref, n) - fused_out = mhc_fused_expand_combine(f, bias, H_post, x, H_res, n, False) + fused_out = mhc_fused_expand_combine(f, bias, H_post, x, H_res, n=n, use_tf32=False) torch.testing.assert_close(fused_out, ref_out, **tols) diff --git a/tests/pytorch/test_mxfp8_2d_quantize.py b/tests/pytorch/test_mxfp8_2d_quantize.py new file mode 100644 index 0000000000..a0159720c5 --- /dev/null +++ b/tests/pytorch/test_mxfp8_2d_quantize.py @@ -0,0 +1,482 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for MXFP8 2D quantization.""" + +import pytest +import torch + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine.common.recipe import MXFP8BlockScaling +from transformer_engine.pytorch import MXFP8Quantizer +from transformer_engine.pytorch.quantization import ( + MXFP8BlockScalingRecipeState, + QuantizerRole, +) + + +mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) +MXFP8_BLOCK_SIZE = 32 +FP8_E4M3_MAX = 448.0 +MXFP8_TEST_SHAPES = [ + (64, 64), + (128, 128), + (256, 1024), + (1024, 256), + (256, 288), + (320, 320), + (352, 256), + (2048, 1024), +] +MXFP8_TEST_DTYPES = [torch.float32, torch.bfloat16] + + +def _valid_rowwise_scale(scale: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """Return the logical, unpadded rowwise scale region.""" + return scale[:rows, : (cols + MXFP8_BLOCK_SIZE - 1) // MXFP8_BLOCK_SIZE] + + +def _valid_columnwise_scale(scale: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """Return the logical, unpadded columnwise scale region.""" + return scale[: (rows + MXFP8_BLOCK_SIZE - 1) // MXFP8_BLOCK_SIZE, :cols] + + +def _float_to_e8m0(amax: torch.Tensor) -> torch.Tensor: + """Convert amax values to E8M0 scale bytes with the same ceil policy as TE.""" + val = (amax.to(torch.float32) / FP8_E4M3_MAX).contiguous() + val_u32 = val.view(torch.int32) + exponent = ((val_u32 >> 23) & 0xFF).to(torch.int32) + mantissa = val_u32 & 0x7FFFFF + + round_up = (mantissa > 0) & (exponent != 254) & ~((exponent == 0) & (mantissa <= 0x400000)) + exponent = exponent + round_up.to(torch.int32) + exponent = torch.where(val == 0, torch.zeros_like(exponent), exponent) + + return exponent.to(torch.uint8) + + +def _e8m0_to_scale_inv(e8m0: torch.Tensor) -> torch.Tensor: + """Convert E8M0 scale bytes back to scale-inverse values.""" + return torch.pow(2.0, e8m0.to(torch.float32) - 127) + + +def _mxfp8_2d_quantize_reference( + x: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Reference MXFP8 2D quantization using one scale per 32x32 block.""" + rows, cols = x.shape + assert rows % MXFP8_BLOCK_SIZE == 0 + assert cols % MXFP8_BLOCK_SIZE == 0 + + block_rows = rows // MXFP8_BLOCK_SIZE + block_cols = cols // MXFP8_BLOCK_SIZE + + x_blocks = x.view( + block_rows, + MXFP8_BLOCK_SIZE, + block_cols, + MXFP8_BLOCK_SIZE, + ).permute(0, 2, 1, 3) + + block_amax = torch.amax(torch.abs(x_blocks.to(torch.float32)), dim=(-1, -2)) + block_scale_e8m0 = _float_to_e8m0(block_amax) + block_scale_inv = _e8m0_to_scale_inv(block_scale_e8m0) + + x_scaled = x_blocks.to(torch.float32) / block_scale_inv[:, :, None, None] + x_quantized = x_scaled.to(torch.float8_e4m3fn) + rowwise_data = x_quantized.permute(0, 2, 1, 3).reshape(rows, cols) + rowwise_scale = block_scale_e8m0.repeat_interleave(MXFP8_BLOCK_SIZE, dim=0) + columnwise_scale = block_scale_e8m0.repeat_interleave(MXFP8_BLOCK_SIZE, dim=1) + + return rowwise_data, rowwise_scale, columnwise_scale + + +def _quantize( + quantizer: MXFP8Quantizer, + x: torch.Tensor, + use_preallocated_output: bool, +) -> torch.Tensor: + """Quantize with either the C++ allocator or an explicitly preallocated output.""" + if not use_preallocated_output: + return quantizer(x) + + out = quantizer.make_empty( + x.shape, + dtype=x.dtype, + device=x.device, + requires_grad=False, + ) + return quantizer.update_quantized(x, out) + + +def _assert_rowwise_scales_are_2d(scales: torch.Tensor, rows: int, cols: int) -> None: + """Check that each 32x32 block uses one rowwise scale for all rows.""" + valid = _valid_rowwise_scale(scales, rows, cols) + block_rows = (rows + MXFP8_BLOCK_SIZE - 1) // MXFP8_BLOCK_SIZE + block_cols = valid.shape[1] + for block_row in range(block_rows): + row_start = block_row * MXFP8_BLOCK_SIZE + row_end = min(row_start + MXFP8_BLOCK_SIZE, rows) + for block_col in range(block_cols): + block_scales = valid[row_start:row_end, block_col] + torch.testing.assert_close( + block_scales, + block_scales[0].expand_as(block_scales), + atol=0, + rtol=0, + ) + + +def _assert_bidirectional_scales_are_2d( + rowwise_scales: torch.Tensor, + columnwise_scales: torch.Tensor, + rows: int, + cols: int, +) -> None: + """Check that rowwise and columnwise metadata agree per 32x32 block.""" + rowwise_valid = _valid_rowwise_scale(rowwise_scales, rows, cols) + columnwise_valid = _valid_columnwise_scale(columnwise_scales, rows, cols) + block_rows = columnwise_valid.shape[0] + block_cols = rowwise_valid.shape[1] + + for block_row in range(block_rows): + row_start = block_row * MXFP8_BLOCK_SIZE + row_end = min(row_start + MXFP8_BLOCK_SIZE, rows) + for block_col in range(block_cols): + col_start = block_col * MXFP8_BLOCK_SIZE + col_end = min(col_start + MXFP8_BLOCK_SIZE, cols) + rowwise_block_scales = rowwise_valid[row_start:row_end, block_col] + columnwise_block_scales = columnwise_valid[block_row, col_start:col_end] + torch.testing.assert_close( + rowwise_block_scales, + rowwise_block_scales[0].expand_as(rowwise_block_scales), + atol=0, + rtol=0, + ) + torch.testing.assert_close( + columnwise_block_scales, + columnwise_block_scales[0].expand_as(columnwise_block_scales), + atol=0, + rtol=0, + ) + torch.testing.assert_close( + rowwise_block_scales[0], + columnwise_block_scales[0], + atol=0, + rtol=0, + ) + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +@pytest.mark.parametrize("columnwise", [False, True], ids=["rowwise_only", "bidirectional"]) +def test_mxfp8_2d_quantize_scales_match_known_block_amax(columnwise: bool) -> None: + """Check exact 2D MXFP8 scale bytes on a hand-built 2x2 block matrix. + + The input has one known amax per 32x32 block. Each amax is chosen so the + expected E8M0 scale byte is known exactly, which makes this test independent + of the random-input comparisons below. Rowwise-only mode should emit only + rowwise tensors, while bidirectional mode should emit matching rowwise and + columnwise scale metadata for the same 2D blocks. + """ + rows, cols = 64, 64 + x = torch.zeros((rows, cols), dtype=torch.float32, device="cuda") + block_exponents = torch.tensor([[-2, -1], [0, 1]], device="cuda") + expected_block_scales = (block_exponents + 127).to(torch.uint8) + + for block_row in range(block_exponents.shape[0]): + for block_col in range(block_exponents.shape[1]): + row_start = block_row * MXFP8_BLOCK_SIZE + col_start = block_col * MXFP8_BLOCK_SIZE + amax = 448.0 * (2.0 ** int(block_exponents[block_row, block_col].item())) + x[ + row_start : row_start + MXFP8_BLOCK_SIZE, + col_start : col_start + MXFP8_BLOCK_SIZE, + ] = ( + amax * 0.5 + ) + x[row_start, col_start] = amax + + quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=columnwise, + with_2d_quantization=True, + ) + out = quantizer(x) + + expected_rowwise_scales = expected_block_scales.repeat_interleave(MXFP8_BLOCK_SIZE, dim=0) + torch.testing.assert_close( + _valid_rowwise_scale(out._rowwise_scale_inv, rows, cols), + expected_rowwise_scales, + atol=0, + rtol=0, + ) + + if columnwise: + expected_columnwise_scales = expected_block_scales.repeat_interleave( + MXFP8_BLOCK_SIZE, dim=1 + ) + torch.testing.assert_close( + _valid_columnwise_scale(out._columnwise_scale_inv, rows, cols), + expected_columnwise_scales, + atol=0, + rtol=0, + ) + else: + assert out._columnwise_data is None + assert out._columnwise_scale_inv is None + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +@pytest.mark.parametrize("shape", MXFP8_TEST_SHAPES) +@pytest.mark.parametrize("dtype", MXFP8_TEST_DTYPES, ids=str) +@pytest.mark.parametrize("columnwise", [False, True], ids=["rowwise_only", "bidirectional"]) +@pytest.mark.parametrize( + "use_preallocated_output", + [False, True], + ids=["cpp_allocator", "preallocated_output"], +) +def test_mxfp8_2d_quantize_matches_torch_reference( + shape: tuple[int, int], + dtype: torch.dtype, + columnwise: bool, + use_preallocated_output: bool, +) -> None: + """Compare random-input MXFP8 2D data and scales against a PyTorch reference.""" + rows, cols = shape + torch.manual_seed(9012) + x = torch.randn(shape, dtype=dtype, device="cuda") + + quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=columnwise, + with_2d_quantization=True, + ) + out = _quantize(quantizer, x, use_preallocated_output) + + ref_data, ref_rowwise_scale, ref_columnwise_scale = _mxfp8_2d_quantize_reference(x) + ref_data_uint8 = ref_data.view(torch.uint8) + + assert out._rowwise_data is not None + assert out._rowwise_scale_inv is not None + torch.testing.assert_close( + out._rowwise_data.view(torch.uint8), + ref_data_uint8, + atol=0, + rtol=0, + ) + torch.testing.assert_close( + _valid_rowwise_scale(out._rowwise_scale_inv, rows, cols), + ref_rowwise_scale, + atol=0, + rtol=0, + ) + + if columnwise: + assert out._columnwise_data is not None + assert out._columnwise_scale_inv is not None + torch.testing.assert_close( + out._columnwise_data.view(torch.uint8), + ref_data_uint8, + atol=0, + rtol=0, + ) + torch.testing.assert_close( + _valid_columnwise_scale(out._columnwise_scale_inv, rows, cols), + ref_columnwise_scale, + atol=0, + rtol=0, + ) + else: + assert out._columnwise_data is None + assert out._columnwise_scale_inv is None + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +@pytest.mark.parametrize("shape", MXFP8_TEST_SHAPES) +@pytest.mark.parametrize("dtype", MXFP8_TEST_DTYPES, ids=str) +@pytest.mark.parametrize( + "use_preallocated_output", + [False, True], + ids=["cpp_allocator", "preallocated_output"], +) +def test_mxfp8_2d_quantize_rowwise_only_matches_bidirectional( + shape: tuple[int, int], + dtype: torch.dtype, + use_preallocated_output: bool, +) -> None: + """2D MXFP8 must support inference-style rowwise-only weight quantization.""" + rows, cols = shape + torch.manual_seed(1234) + x = torch.randn(shape, dtype=dtype, device="cuda") + + rowwise_only_quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + with_2d_quantization=True, + ) + bidirectional_quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + with_2d_quantization=True, + ) + + rowwise_only = _quantize(rowwise_only_quantizer, x, use_preallocated_output) + bidirectional = _quantize(bidirectional_quantizer, x, use_preallocated_output) + + assert rowwise_only._rowwise_data is not None + assert rowwise_only._rowwise_scale_inv is not None + assert rowwise_only._columnwise_data is None + assert rowwise_only._columnwise_scale_inv is None + + torch.testing.assert_close( + rowwise_only._rowwise_data.view(torch.uint8), + bidirectional._rowwise_data.view(torch.uint8), + atol=0, + rtol=0, + ) + torch.testing.assert_close( + _valid_rowwise_scale(rowwise_only._rowwise_scale_inv, rows, cols), + _valid_rowwise_scale(bidirectional._rowwise_scale_inv, rows, cols), + atol=0, + rtol=0, + ) + _assert_rowwise_scales_are_2d(rowwise_only._rowwise_scale_inv, rows, cols) + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +@pytest.mark.parametrize("shape", MXFP8_TEST_SHAPES) +@pytest.mark.parametrize("dtype", MXFP8_TEST_DTYPES, ids=str) +@pytest.mark.parametrize( + "use_preallocated_output", + [False, True], + ids=["cpp_allocator", "preallocated_output"], +) +def test_mxfp8_2d_quantize_bidirectional_scales_match( + shape: tuple[int, int], + dtype: torch.dtype, + use_preallocated_output: bool, +) -> None: + """Rowwise and columnwise scale metadata should encode the same 32x32 block scales.""" + rows, cols = shape + torch.manual_seed(5678) + x = torch.randn(shape, dtype=dtype, device="cuda") + + quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + with_2d_quantization=True, + ) + out = _quantize(quantizer, x, use_preallocated_output) + + assert out._rowwise_scale_inv is not None + assert out._columnwise_scale_inv is not None + _assert_bidirectional_scales_are_2d( + out._rowwise_scale_inv, + out._columnwise_scale_inv, + rows, + cols, + ) + + +def test_mxfp8_recipe_state_uses_2d_only_for_forward_weights() -> None: + """Only forward weight quantizers should inherit MXFP8 2D quantization.""" + recipe = MXFP8BlockScaling(enable_2d_quantization=True) + roles = [ + QuantizerRole(module_type="linear", tensor_type="input"), + QuantizerRole(module_type="linear", tensor_type="weight"), + QuantizerRole(module_type="linear", tensor_type="output"), + ] + state = MXFP8BlockScalingRecipeState( + recipe=recipe, mode="forward", num_quantizers=3, roles=roles + ) + quantizers = state.make_quantizers() + + assert [q.with_2d_quantization for q in quantizers] == [False, True, False] + + backward_state = MXFP8BlockScalingRecipeState( + recipe=recipe, + mode="backward", + num_quantizers=2, + roles=[ + QuantizerRole(module_type="linear", tensor_type="grad_output"), + QuantizerRole(module_type="linear", tensor_type="grad_input"), + ], + ) + assert [q.with_2d_quantization for q in backward_state.make_quantizers()] == [ + False, + False, + ] + + +def test_mxfp8_recipe_state_2d_requires_explicit_weight_role() -> None: + """MXFP8 2D should not enable itself for unknown positional slots.""" + recipe = MXFP8BlockScaling(enable_2d_quantization=True) + state = MXFP8BlockScalingRecipeState(recipe=recipe, mode="forward", num_quantizers=3) + assert [q.with_2d_quantization for q in state.make_quantizers()] == [ + False, + False, + False, + ] + + +def test_mxfp8_recipe_state_uses_2d_for_grouped_linear_weights() -> None: + """GroupedLinear weight quantizers should inherit MXFP8 2D quantization.""" + recipe = MXFP8BlockScaling(enable_2d_quantization=True) + roles = [ + QuantizerRole(module_type="grouped_linear", tensor_type="input"), + QuantizerRole(module_type="grouped_linear", tensor_type="weight"), + QuantizerRole(module_type="grouped_linear", tensor_type="output"), + ] + state = MXFP8BlockScalingRecipeState( + recipe=recipe, + mode="forward", + num_quantizers=len(roles), + roles=roles, + ) + assert [q.with_2d_quantization for q in state.make_quantizers()] == [ + False, + True, + False, + ] + + +def test_mxfp8_recipe_state_2d_ignores_unsupported_roles() -> None: + """MXFP8 2D is limited to supported Linear weight quantizers.""" + recipe = MXFP8BlockScaling(enable_2d_quantization=True) + roles = [ + QuantizerRole(module_type="dpa", tensor_type="qkv"), + QuantizerRole(module_type="dpa", tensor_type="weight"), + QuantizerRole(module_type="", tensor_type="weight"), + ] + state = MXFP8BlockScalingRecipeState( + recipe=recipe, + mode="forward", + num_quantizers=len(roles), + roles=roles, + ) + assert [q.with_2d_quantization for q in state.make_quantizers()] == [ + False, + False, + False, + ] + + +def test_mxfp8_quantizer_copy_preserves_2d_flag() -> None: + """MXFP8Quantizer.copy should preserve the 2D quantization setting.""" + quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + with_2d_quantization=True, + ) + copied = quantizer.copy() + assert copied.rowwise_usage is True + assert copied.columnwise_usage is False + assert copied.with_2d_quantization is True diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 71fea22b66..6cdd784216 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -49,6 +49,10 @@ ) from transformer_engine.pytorch.attention.dot_product_attention.utils import FlashAttentionUtils as fa_utils from transformer_engine.pytorch import checkpoint as te_checkpoint +from transformer_engine.pytorch.distributed import ( + is_fp8_activation_recompute_enabled, + in_fp8_activation_recompute_phase, +) from transformer_engine.pytorch.cpp_extensions import general_gemm from transformer_engine.common import recipe from transformer_engine.pytorch import DType @@ -100,6 +104,9 @@ def rocm_attn_backend() -> tuple[bool, bool, bool]: all_boolean = [True, False] +# fp8_meta key written by FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute +_FP8_RECOMPUTE_KEY = "global_fp8_buffer_pos_fwd_recompute" + all_activations = [ "gelu", "geglu", @@ -680,7 +687,15 @@ def test_gpt_selective_activation_recompute(dtype, bs, model, fp8, recipe, fp8_m def _test_e2e_full_recompute( - bs, dtype, config, fp8, recipe, fp8_model_params=False, recompute=False, use_reentrant=True + bs, + dtype, + config, + fp8, + recipe, + fp8_model_params=False, + recompute=False, + use_reentrant=True, + inner_autocast=False, ): reset_rng_states() FP8GlobalStateManager.reset() @@ -717,10 +732,17 @@ def _test_e2e_full_recompute( te_inp_hidden_states.retain_grad() te_inp_attn_mask = get_causal_attn_mask(config.max_seqlen_q) - with autocast(enabled=fp8, recipe=recipe): + forward = block + if inner_autocast: + + def forward(*args, **kwargs): + with autocast(enabled=fp8, recipe=recipe): + return block(*args, **kwargs) + + with autocast(enabled=fp8 and not inner_autocast, recipe=recipe): if recompute: te_out = te_checkpoint( - block, + forward, te_inp_hidden_states, attention_mask=te_inp_attn_mask, checkpoint_core_attention=False, @@ -729,7 +751,7 @@ def _test_e2e_full_recompute( use_reentrant=use_reentrant, ) else: - te_out = block( + te_out = forward( te_inp_hidden_states, attention_mask=te_inp_attn_mask, checkpoint_core_attention=False, @@ -830,6 +852,152 @@ def test_gpt_full_activation_recompute( ) +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("use_reentrant", all_boolean) +def test_gpt_full_activation_recompute_with_inner_autocast(use_reentrant, monkeypatch): + """Check recompute numerics when FP8 autocast starts inside the checkpointed callable.""" + if not use_reentrant: + # Non-reentrant checkpoint becomes non-deterministic with bias+GELU fusion. + monkeypatch.setenv("NVTE_BIAS_GELU_NVFUSION", "0") + + dtype = torch.bfloat16 + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) + config = model_configs["126m"] + + # Reference also opens the autocast inside the callable, so the only difference + # between the two runs is activation recompute. + outputs, names = _test_e2e_full_recompute( + 1, + dtype, + config, + True, + fp8_recipe, + recompute=False, + use_reentrant=use_reentrant, + inner_autocast=True, + ) + + # Before the fix, phase 1 skipped the stash while the recompute still restored it, which + # surfaced as a KeyError on the recompute buffer lookup. Count both sides to pin that down. + stash_counts, restore_counts = {}, {} + stash_fn = FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute + restore_fn = FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute + + def record_stash(fp8_meta): + stash_fn(fp8_meta) + if _FP8_RECOMPUTE_KEY in fp8_meta: + stash_counts[id(fp8_meta)] = stash_counts.get(id(fp8_meta), 0) + 1 + + def record_restore(fp8_meta): + # The restore site is not gated on delayed scaling, but only delayed scaling stashes. + if not fp8_meta["recipe"].delayed(): + restore_fn(fp8_meta) + return + key = id(fp8_meta) + assert key in stash_counts, "Recompute restored a scale that was never stashed" + restore_counts[key] = restore_counts.get(key, 0) + 1 + restore_fn(fp8_meta) + + monkeypatch.setattr( + FP8GlobalStateManager, + "copy_forward_fp8_meta_tensors_for_recompute", + staticmethod(record_stash), + ) + monkeypatch.setattr( + FP8GlobalStateManager, + "get_old_fp8_meta_tensors_for_recompute", + staticmethod(record_restore), + ) + + outputs_recompute, _ = _test_e2e_full_recompute( + 1, + dtype, + config, + True, + fp8_recipe, + recompute=True, + use_reentrant=use_reentrant, + inner_autocast=True, + ) + + assert stash_counts, "No FP8 module stashed a forward scale for the recompute phase" + assert restore_counts == stash_counts, "Stash and restore of forward scales are unbalanced" + + for name, ref, test in zip(names, outputs, outputs_recompute): + torch.testing.assert_close( + test, + ref, + msg=f"Mismatch in tensor {name}", + rtol=0.125, + atol=0.0675, + ) + + +def _checkpointed_linear_backward(body, use_reentrant, *layers): + """Run a checkpointed callable end to end and check the gradients are finite.""" + inp = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True) + with torch.autocast("cuda", dtype=torch.bfloat16): + out = te_checkpoint(body, inp, use_reentrant=use_reentrant) + loss = out.float().sum() + loss.backward() + torch.cuda.synchronize() + + assert torch.isfinite(loss) + assert inp.grad is not None and torch.isfinite(inp.grad).all() + for layer in layers: + assert layer.weight.grad is not None + assert torch.isfinite(layer.weight.grad).all() + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("use_reentrant", all_boolean) +def test_checkpoint_inner_autocast_is_an_fp8_recompute_region(use_reentrant): + """An FP8 autocast opened inside a checkpointed callable is an FP8 recompute region.""" + FP8GlobalStateManager.reset() + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) + layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda() + + observed = [] + + def body(value): + outside = is_fp8_activation_recompute_enabled() + with autocast(enabled=True, recipe=fp8_recipe): + observed.append( + ( + outside, + is_fp8_activation_recompute_enabled(), + in_fp8_activation_recompute_phase(), + ) + ) + return layer(value) + + _checkpointed_linear_backward(body, use_reentrant, layer) + + # One entry for the checkpointed forward, one for the recompute during backward. The + # query is only an FP8 recompute region inside the autocast, in both phases. + assert observed == [(False, True, False), (False, True, True)] + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("use_reentrant", all_boolean) +def test_checkpoint_with_mixed_fp8_regions_saves_only_fp8_recompute_state(use_reentrant): + """Only the inner FP8 region of a mixed checkpoint saves recompute metadata.""" + FP8GlobalStateManager.reset() + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) + non_fp8_layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda() + fp8_layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda() + + def body(value): + value = non_fp8_layer(value) + with autocast(enabled=True, recipe=fp8_recipe): + return fp8_layer(value) + + _checkpointed_linear_backward(body, use_reentrant, non_fp8_layer, fp8_layer) + + assert _FP8_RECOMPUTE_KEY not in non_fp8_layer.fp8_meta + assert _FP8_RECOMPUTE_KEY in fp8_layer.fp8_meta + + def _test_e2e_checkpointing_get_model(config, dtype): sigma = 0.023 init_method = init_method_normal(sigma) diff --git a/tests/pytorch/test_onnx_export.py b/tests/pytorch/test_onnx_export.py index 0ba1536882..6ddf9cbbbc 100644 --- a/tests/pytorch/test_onnx_export.py +++ b/tests/pytorch/test_onnx_export.py @@ -368,6 +368,11 @@ def create_ort_input_dict(session, inputs): ) +def get_atol(precision: torch.dtype, default: float = 1e-3) -> float: + # ORT runs on CPU while TE runs on GPU; fp16 accumulation differences reach a few ULPs. + return 2e-2 if precision is torch.float16 else default + + def dtype2str(dtype: torch.dtype, fake_bf16_io=False): if fake_bf16_io: assert dtype == torch.bfloat16 @@ -458,7 +463,7 @@ def forward(self, inp): if precision in (torch.bfloat16,): return if fp8_recipe is None: - validate_result(fname, inp, model, atol=1e-3, te_outputs=te_outputs) + validate_result(fname, inp, model, atol=get_atol(precision), te_outputs=te_outputs) else: validate_result( fname, inp, model, atol=1e-2, is_fp8=fp8_recipe is not None, te_outputs=te_outputs @@ -587,7 +592,7 @@ def _test_export_layernorm_linear( if precision in (torch.bfloat16,): return if fp8_recipe is None: - validate_result(fname, inp, model, atol=1e-3, te_outputs=te_outputs) + validate_result(fname, inp, model, atol=get_atol(precision), te_outputs=te_outputs) elif precision != torch.bfloat16: validate_result( fname, @@ -595,7 +600,11 @@ def _test_export_layernorm_linear( model, # For current scaling we use Float8Quantizer in tests + amax computed by hand, # which has slightly different numerics than Float8CurrentScalingQuantizer. - atol=1e-3 if fp8_recipe.__class__ is not recipe.Float8CurrentScaling else 2e-2, + atol=( + get_atol(precision) + if fp8_recipe.__class__ is not recipe.Float8CurrentScaling + else 2e-2 + ), is_fp8=fp8_recipe is not None, te_outputs=te_outputs, ) @@ -673,7 +682,9 @@ def _test_export_layernorm_mlp( if precision in (torch.bfloat16,): return atol = ( - 2e-2 if fp8_recipe is not None else (5e-1 if activation == "swiglu" else 1e-3) + 2e-2 + if fp8_recipe is not None + else (5e-1 if activation == "swiglu" else get_atol(precision)) ) # TODO(pgadzinski) - check 2e-2 validate_result( fname, inp, model, atol=atol, is_fp8=fp8_recipe is not None, te_outputs=te_outputs diff --git a/tests/pytorch/test_ops_grouped_linear_distributed_weight.py b/tests/pytorch/test_ops_grouped_linear_distributed_weight.py new file mode 100644 index 0000000000..700be4138e --- /dev/null +++ b/tests/pytorch/test_ops_grouped_linear_distributed_weight.py @@ -0,0 +1,131 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""DistributedWeight dispatch in the fusible ``ops.GroupedLinear`` (unfused fallback path). + +TE ships no DistributedWeight implementer (GTP etc. live in the caller), so this validates the +*dispatch wiring* with an in-repo fake. The fake applies a DISTINCT, observable scale in each +materialize hook so every plumbing point is independently checked -- a wiring bug fails a specific +assertion: + + * ``materialize_group_for_forward`` scales the weights by ``FWD_SCALE`` -> the fwd GEMM must use + the materialized weights, so ``out == FWD_SCALE * plain_out``. + * ``materialize_group_for_backward`` scales by ``BWD_SCALE`` -> dgrad must use the + re-materialized weights, so ``dgrad == BWD_SCALE * plain_dgrad``. + * ``finalize_group_grads`` reduce-scatters the wgrads into ``main_grad`` in-place and returns a + dummy -> the real wgrad lands in ``main_grad`` and the ops path returns a throwaway ``.grad`` + (it discards finalize's return; see DistributedWeight.finalize_group_grads). + +Real all-gather / reduce-scatter math is exercised by the caller's distributed tests; here the fake +is single-process and only proves the ops.GroupedLinear integration routes weights/grads through +the hooks (and does not bypass them). +""" + +import pytest +import torch + +import transformer_engine.pytorch as te + +# Distinct powers of two so each scale commutes exactly through the (possibly TF32) GEMM rounding, +# making the linear scale relations bit-exact under a tight tolerance. +FWD_SCALE = 2.0 +BWD_SCALE = 4.0 + + +class _FakeDistWeight(torch.nn.Parameter): + """Single-process fake DistributedWeight leader for dispatch testing. + + Each materialize hook applies its own scale (see module docstring) so fwd/bwd routing is + observable; ``finalize_group_grads`` models the real main-grad contract -- reduce-scatter (here + an identity accumulate) the wgrads into each shard's ``main_grad`` in-place, flag + ``grad_added_to_main_grad``, and return a dummy that the ops path discards. + """ + + is_distributed_weight = True + + def materialize_group_for_forward(self): + self.calls["fwd"] += 1 + return [w * FWD_SCALE for w in self._group] + + def materialize_group_for_backward(self, **kwargs): + self.calls["bwd"] += 1 + return [w * BWD_SCALE for w in self._group] + + def finalize_group_grads(self, wgrads, **kwargs): + self.calls["finalize"] += 1 + wl = list(wgrads) if isinstance(wgrads, (list, tuple)) else [wgrads] + for w, g in zip(self._group, wl): + w.main_grad.add_(g.to(w.main_grad.dtype)) # in-place accumulate into main_grad + w.grad_added_to_main_grad = True + return [torch.zeros_like(g) for g in wl] # dummy grads (real value is now in main_grad) + + def grad_buffer(self): + return self.data + + +def _make_fake_dist_leader(op, num_gemms): + """Replace ``op.weight0`` with a fake distributed leader referencing the whole group.""" + w0 = op.weight0 + leader = _FakeDistWeight(w0.data) + leader.calls = {"fwd": 0, "bwd": 0, "finalize": 0} + leader._group = [leader] + [getattr(op, f"weight{i}") for i in range(1, num_gemms)] + op.weight0 = leader + return leader + + +@pytest.mark.parametrize("num_gemms", [2, 4]) +def test_ops_grouped_linear_distributed_weight_dispatch(num_gemms): + """Every DistributedWeight hook must be routed through the GEMM flow (and not bypassed). + + fwd/bwd use the scaled materialized weights; finalize reduce-scatters the wgrad into + ``main_grad`` in-place, and the ops path returns a throwaway dummy ``.grad``. + """ + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.manual_seed(0) + # fp32 with power-of-two scales: each scale commutes exactly through the GEMM rounding (even + # TF32), so the linear scale relations below are bit-exact under a tight tolerance. + in_f, out_f, total_tokens = 32, 64, num_gemms * 8 + dtype, device = torch.float32, "cuda" + + op = te.ops.GroupedLinear(num_gemms, in_f, out_f, bias=False, device=device, dtype=dtype) + reference = te.ops.GroupedLinear(num_gemms, in_f, out_f, bias=False, device=device, dtype=dtype) + reference.load_state_dict(op.state_dict()) + + leader = _make_fake_dist_leader(op, num_gemms) + for i in range(num_gemms): + w = getattr(op, f"weight{i}") + w.main_grad = torch.zeros((out_f, in_f), dtype=torch.float32, device=device) + w.grad_added_to_main_grad = False # DDP initializes this on every param + + m_splits = [total_tokens // num_gemms] * num_gemms + m_splits[-1] += total_tokens - sum(m_splits) + split_sizes = torch.tensor(m_splits, dtype=torch.int64, device=device) + + x = torch.randn(total_tokens, in_f, dtype=dtype, device=device, requires_grad=True) + ref_x = x.detach().clone().requires_grad_(True) + + out = op(x, split_sizes) + out.sum().backward() + ref_out = reference(ref_x, split_sizes) + ref_out.sum().backward() + + # All three dispatch hooks actually fired. + assert leader.calls["fwd"] > 0 and leader.calls["bwd"] > 0 and leader.calls["finalize"] > 0 + + tols = dict(rtol=1e-5, atol=1e-5) + # fwd used the materialized (FWD_SCALE) weights. + torch.testing.assert_close(out, FWD_SCALE * ref_out, **tols) + # dgrad used the re-materialized (BWD_SCALE) weights. + torch.testing.assert_close(x.grad, BWD_SCALE * ref_x.grad, **tols) + for i in range(num_gemms): + w = getattr(op, f"weight{i}") + ref_w = getattr(reference, f"weight{i}") + # The distributed op's grad of record is main_grad (finalize reduce-scattered the identity + # wgrad there); compare it to the plain reference module's ordinary autograd .grad. + torch.testing.assert_close(w.main_grad.to(dtype), ref_w.grad, **tols) + # main_grad is flagged, and op's own .grad is a discarded dummy (not the real grad, which + # lives in main_grad). Nobody writes the real grad into op.weight.grad by design. + assert w.grad_added_to_main_grad is True + assert w.grad is not None, f"weight{i} should receive a dummy .grad" diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index b2f77ecd66..4bcaacac90 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/tests/pytorch/test_quantized_tensor.py @@ -20,6 +20,7 @@ Float8Tensor, Float8BlockwiseQTensor, MXFP8Tensor, + MXFP8TensorStorage, NVFP4Tensor, QuantizedTensor, ) @@ -670,6 +671,31 @@ def test_cpu_dequantize( assert y_cpu.shape == ref_cpu.shape torch.testing.assert_close(y_cpu, ref_cpu, rtol=0, atol=0) + def test_mxfp8_fsdp_extract_keeps_partial_columnwise_scale_block(self) -> None: + """FSDP extraction must retain the scale for a partial final 32-row block.""" + rows, cols = 48, 64 + columnwise_scale_inv = torch.arange(4 * 128, dtype=torch.float32).reshape(4, 128) + storage = MXFP8TensorStorage( + rowwise_data=None, + rowwise_scale_inv=None, + columnwise_data=torch.zeros((rows, cols), dtype=torch.uint8), + columnwise_scale_inv=columnwise_scale_inv, + fp8_dtype=te.DType.kFloat8E4M3, + quantizer=None, + with_gemm_swizzled_scales=False, + fake_dtype=torch.bfloat16, + ) + + buffers, metadata = storage.fsdp_extract_buffers() + + assert metadata == {"field_names": ("_columnwise_data", "_columnwise_scale_inv")} + assert len(buffers) == 2 + assert buffers[0] is storage._columnwise_data + extracted_scale_inv = buffers[1] + assert extracted_scale_inv is not None + assert extracted_scale_inv.shape == (2, 128) + torch.testing.assert_close(extracted_scale_inv, columnwise_scale_inv[:2], rtol=0, atol=0) + @pytest.mark.parametrize("quantization", _quantization_list) @pytest.mark.parametrize("dim", [0, 1]) def test_chunk( @@ -713,6 +739,14 @@ def test_chunk( y_test = y_test.to(dtype=torch.float64, device="cpu") torch.testing.assert_close(y_test, y_ref, **tols) + def test_view_not_implemented(self) -> None: + """QuantizedTensor base class does not support tensor views.""" + qt = QuantizedTensor((128, 128), torch.bfloat16) + with pytest.raises( + NotImplementedError, match="QuantizedTensor class does not support tensor views" + ): + qt.view(-1) + @pytest.mark.parametrize("quantization", _quantization_list) def test_shape_with_none_data( self, @@ -757,6 +791,58 @@ def test_shape_with_none_data( f"after setting data to None on {type(x_test).__name__}" ) + @pytest.mark.parametrize("quantization", _quantization_list) + @pytest.mark.parametrize( + "rowwise, columnwise", + [(True, True), (True, False), (False, True)], + ids=["rowwise_columnwise", "rowwise_only", "columnwise_only"], + ) + @pytest.mark.parametrize("shape", [(128, 256), (4, 128, 256)], ids=["2d", "3d"]) + def test_shape_matches_size( + self, + *, + quantization: str, + rowwise: bool, + columnwise: bool, + shape: Iterable[int], + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + ) -> None: + """shape, size() and size(dim) stay consistent for every usage combination. + + Both shape and size() are derived from whichever data buffer is present, + and classes that store columnwise data transposed have to undo that. A + columnwise-only tensor is where they can drift apart -- from each other, + and from the shape the tensor was allocated with. + """ + quantizer = make_quantizer(quantization, device=device) + # Row-scaled NVFP4 accepts set_usage(rowwise=False) but rejects the + # allocation itself, so it has to be filtered out up front. + if getattr(quantizer, "row_scaled_nvfp4", False) and not rowwise: + pytest.skip(f"{quantization} requires rowwise usage") + quantizer.set_usage(rowwise=rowwise, columnwise=columnwise) + if (quantizer.rowwise_usage, quantizer.columnwise_usage) != (rowwise, columnwise): + pytest.skip(f"{quantization} does not support this usage combination") + + x = quantizer.make_empty(shape, dtype=dtype, device=device) + name = type(x).__name__ + + # shape and size() must describe the same tensor, whichever buffer they + # end up reading. + assert tuple(x.shape) == tuple(x.size()), f"{name}: {tuple(x.shape)} vs {tuple(x.size())}" + + # size(dim) must agree with the full shape, including negative indices. + # It cannot be served by forwarding dim to a transposed buffer. + for dim in range(len(x.shape)): + assert x.size(dim) == x.shape[dim], f"{name}.size({dim}) is {x.size(dim)}" + neg = dim - len(x.shape) + assert x.size(neg) == x.shape[neg], f"{name}.size({neg}) is {x.size(neg)}" + + # NVFP4 deliberately reports columnwise-only tensors flattened to 2D and + # warns about it, so only the ranks it preserves are checked here. + if not (isinstance(quantizer, NVFP4Quantizer) and not rowwise and len(shape) > 2): + assert tuple(x.shape) == tuple(shape), f"{name}.shape is {tuple(x.shape)}" + @pytest.mark.parametrize( "quantization", _quantization_list + (["nvfp4_2d"] if nvfp4_available else []), @@ -928,3 +1014,96 @@ def test_mxfp8_dequantize_columnwise_only_quantized_separately( # Make sure we are not trivially passing the test with pytest.raises(AssertionError): torch.testing.assert_close(x_deq, -x_ref, **_tols[fp8_dtype]) + + +@pytest.mark.parametrize("quantization", _quantization_list) +@pytest.mark.parametrize("usage", ["rowwise", "columnwise", "both"]) +@pytest.mark.parametrize("shape", [(256, 512), (128, 320)], ids=lambda s: f"{s[0]}x{s[1]}") +def test_skip_quantization_with_noop_flag( + quantization: str, usage: str, shape: Tuple[int, int] +) -> None: + """ + Test if the quantization honors the noop flag and skips quantization. + This test only verifies that the kernel skips quantization where it's expected to do so. + It doesn't verify otherwise (kernel correctly quantizes when noop is false) since that is already covered by other tests. + """ + if usage == "columnwise" and quantization in ("fp8", "fp8_delayed_scaling"): + pytest.skip("Delayed scaling does not support columnwise-only quantization") + if usage != "rowwise" and quantization == "nvfp4_row_scaled": + pytest.skip("Row-scaled NVFP4 does not produce columnwise output") + + quantizer = make_quantizer(quantization) + quantizer.set_usage( + rowwise=usage in ("rowwise", "both"), + columnwise=usage in ("columnwise", "both"), + ) + + first_batch = torch.rand(shape, dtype=torch.bfloat16, device="cuda") + # Make the input range much larger, even under delayed scaling, so amax will differ for sure + # if the noop flag doesn't take effect + second_batch = torch.rand_like(first_batch) * 2.0 + x = first_batch.clone() + + noop = torch.zeros(1, dtype=torch.float32, device="cuda") + + # Allocate and populate the destination outside the graph + quantized = quantizer(x) + + # CUDA graph capture requires the work to be warmed up on a side stream first. + side_stream = torch.cuda.Stream() + side_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side_stream): + for _ in range(3): + quantized.quantize_(x, noop_flag=noop) + torch.cuda.current_stream().wait_stream(side_stream) + + # Capture the CUDA graph + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + quantized.quantize_(x, noop_flag=noop) + + # Clone the underlying buffers of the quantized tensor's output so we can compare them later. + # Note that we can't clone the quantized tensor, but its underlying buffers are torch tensors and can be cloned. + def clone_quantized_output(): + output_buffers = {} + for attr in ( + "_data", + "_transpose", + "_rowwise_data", + "_columnwise_data", + "_scale_inv", + "_rowwise_scale_inv", + "_columnwise_scale_inv", + "_amax_rowwise", + "_amax_columnwise", + ): + buf = getattr(quantized, attr, None) + if buf is not None: + output_buffers[attr] = buf.clone() + # Delayed scaling keeps its global amax on the quantizer rather than on the tensor. + quantizer_amax = getattr(quantizer, "amax", None) + if quantizer_amax is not None: + output_buffers["quantizer.amax"] = quantizer_amax.clone() + assert output_buffers, f"{quantization}: quantized tensor exposes no data buffer" + return output_buffers + + # Obtain the quantized x by setting the noop flag to false (zero) + noop.zero_() + graph.replay() + torch.cuda.synchronize() + quantized_without_noop = clone_quantized_output() + + # Reset x to different values and set the noop flag to true (one). + # The quantization should be skipped so the output should not be different. + x.copy_(second_batch) + noop.fill_(1.0) + graph.replay() + torch.cuda.synchronize() + quantized_with_noop = clone_quantized_output() + + for attr in quantized_without_noop: + buf_without_noop = quantized_without_noop[attr] + buf_with_noop = quantized_with_noop[attr] + assert torch.equal( + buf_without_noop, buf_with_noop + ), f"{quantization}/{usage}: noop flag fails to take effect because {attr} changed." diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index 19c495da96..df4797769c 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -40,6 +40,7 @@ from transformer_engine.common import recipe from transformer_engine.pytorch.cpp_extensions import general_gemm from transformer_engine.pytorch.tensor.utils import replace_raw_data +from transformer_engine.pytorch.module import is_module_grouped_tensor_path_supported from utils import ModelConfig, recipe_id, skip_unsupported_backward_override # Only run FP8 tests on supported devices. @@ -89,11 +90,11 @@ def is_fp8_supported(config: ModelConfig): def nvfp4_vanilla(): - nvfp4_recipe = recipe.NVFP4BlockScaling() - nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() - nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams() - nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams() - return nvfp4_recipe + return recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + disable_2d_quantization=True, + ) def nvfp4_row_scaled(): @@ -126,7 +127,7 @@ def nvfp4_4over6(): if mxfp8_available: fp8_recipes.append(recipe.MXFP8BlockScaling()) if nvfp4_available: - fp8_recipes.append(nvfp4_vanilla()) # TODO: fix check for this + fp8_recipes.append(nvfp4_vanilla()) fp8_recipes.append(nvfp4_4over6()) if fp8_block_scaling_available: fp8_recipes.append(recipe.Float8BlockScaling()) @@ -566,6 +567,7 @@ def test_sanity_linear_with_zero_tokens( out = te_linear(inp_hidden_states) loss = out.sum() loss.backward() + torch.cuda.synchronize() assert out.shape == (num_tokens, ffn_hidden_size) @@ -606,17 +608,34 @@ def test_sanity_grouped_linear( if fp8_recipe is not None: fp8_recipe = copy.deepcopy(fp8_recipe) fp8_recipe.backward_override = backward_override + if single_param and not is_module_grouped_tensor_path_supported( + fp8_recipe, + dtype, + ): + pytest.skip("Single grouped parameters require the native grouped-tensor path") + if single_param: + # Single grouped parameters intentionally have no split-quantize fallback, so this + # test must satisfy the native grouped kernels' shape contract. MCore pads each + # expert's token count to 256; TE requires at least 128-row alignment. Weight K must + # be 64-aligned. + tokens_per_nonempty_expert = bs * config.max_seqlen_q + if tokens_per_nonempty_expert % 128 != 0: + pytest.skip("Single grouped parameters require each nonempty m_split to be 128-aligned") + k_alignment = 64 + if config.hidden_size % k_alignment != 0: + pytest.skip(f"Single grouped parameters require GEMM K to be {k_alignment}-aligned") if fp8_recipe is not None: if not is_fp8_supported(config): pytest.skip("Model config does not support FP8") if fp8_recipe.nvfp4(): - if not getattr(fp8_recipe, "row_scaled_activation", False): - pytest.skip("NVFP4 not supported for grouped linear") - if single_param: - pytest.skip("Row-scaled NVFP4 does not support GroupedTensor grouped linear") - if dtype == torch.float16: - pytest.skip("FP16 output for NVFP4 not supported") + if dtype != torch.bfloat16: + pytest.skip("NVFP4 GroupedLinear requires BF16") + if single_param and not fp8_model_params: + pytest.skip( + "NVFP4 single grouped BF16 primary weights require unsupported non-RHT " + "grouped weight quantization; enable quantized model initialization" + ) use_fp8 = fp8_recipe is not None with quantized_model_init(enabled=use_fp8 and fp8_model_params, recipe=fp8_recipe): @@ -628,6 +647,7 @@ def test_sanity_grouped_linear( params_dtype=dtype, single_grouped_weight=single_param, single_grouped_bias=single_param, + use_grouped_tensor=single_param, ).cuda() # Verify grouped linear exposes a single grouped weight parameter(and bias when applicable). @@ -647,11 +667,25 @@ def test_sanity_grouped_linear( m_splits[-1] = 0 elif empty_split == "middle": m_splits[num_gemms // 2] = 0 + if single_param: + m_splits = torch.tensor(m_splits, dtype=torch.int64, device="cuda") + + if NVTE_TEST_NVINSPECT_ENABLED and single_param: + # DebugQuantizer operates on per-GEMM tensors, while single grouped parameters + # intentionally have no split-quantize fallback. + with pytest.raises( + RuntimeError, + match="TE debug features do not support single grouped parameters", + ): + with autocast(enabled=use_fp8, recipe=fp8_recipe): + te_grouped_linear(inp_hidden_states, m_splits) + return with autocast(enabled=use_fp8, recipe=fp8_recipe): out = te_grouped_linear(inp_hidden_states, m_splits) loss = out.sum() loss.backward() + torch.cuda.synchronize() assert out.shape == (num_tokens, ffn_hidden_size) @@ -1157,6 +1191,120 @@ def test_quantized_model_init_high_precision_init_val(): ), "clear_high_precision_init_val() not work" +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("move", ["cuda", "cpu", "half"]) +def test_quantized_param_attrs_survive_apply(move): + """Attributes attached to a quantized parameter survive nn.Module._apply. + + Quantized parameters implement the flatten protocol, so ``_apply`` moves them + with ``swap_tensors``, which exchanges the parameter's whole ``__dict__``. + Anything attached from the outside rides out on the discarded tensor unless + the module re-attaches it. + """ + with quantized_model_init(preserve_high_precision_init_val=True): + model = Linear(64, 64) + + weight = model.weight + expected = weight.get_high_precision_init_val() + weight.probe_attr = "attached-from-outside" + + if move == "cuda": + model = model.cuda() + elif move == "cpu": + model = model.cpu() + else: + model = model.half() + + weight = model.weight + assert hasattr(weight, "get_high_precision_init_val"), f"accessor lost by .{move}()" + assert hasattr(weight, "clear_high_precision_init_val"), f"accessor lost by .{move}()" + torch.testing.assert_close(weight.get_high_precision_init_val(), expected, rtol=0, atol=0) + assert weight.probe_attr == "attached-from-outside", f"custom attr lost by .{move}()" + + # The accessor must read the surviving parameter, not the discarded one. + weight.clear_high_precision_init_val() + assert weight.get_high_precision_init_val() is None + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +def test_grouped_linear_single_param_preserves_high_precision_init(monkeypatch): + """Grouped MXFP8 and discrete weights produce identical FP32 master initialization.""" + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + num_gemms = 3 + + def make_module(single_grouped_weight): + torch.manual_seed(1234) + torch.cuda.manual_seed(1234) + with quantized_model_init( + enabled=True, + recipe=recipe.MXFP8BlockScaling(), + preserve_high_precision_init_val=True, + ): + return GroupedLinear( + num_gemms=num_gemms, + in_features=32, + out_features=64, + bias=False, + params_dtype=torch.bfloat16, + single_grouped_weight=single_grouped_weight, + ).cuda() + + discrete_module = make_module(False) + grouped_module = make_module(True) + discrete_init_vals = [ + getattr(discrete_module, f"weight{i}").get_high_precision_init_val() + for i in range(num_gemms) + ] + expected_grouped_init = torch.stack(discrete_init_vals, dim=0) + + grouped_weight = grouped_module.weight + assert hasattr(grouped_weight, "get_high_precision_init_val") + assert hasattr(grouped_weight, "clear_high_precision_init_val") + grouped_init = grouped_weight.get_high_precision_init_val() + assert grouped_init.device.type == "cpu" + assert grouped_init.shape == grouped_weight.shape + torch.testing.assert_close( + grouped_init, + expected_grouped_init, + rtol=0, + atol=0, + ) + + # This is the exact layout consumed when constructing the FP32 optimizer master. + discrete_master = expected_grouped_init.float() + grouped_master = grouped_init.float() + torch.testing.assert_close(grouped_master, discrete_master, rtol=0, atol=0) + + grouped_weight.clear_high_precision_init_val() + assert grouped_weight.get_high_precision_init_val() is None + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +def test_grouped_linear_rejects_partial_high_precision_init(monkeypatch): + """Packing fails rather than mixing preserved and dequantized initialization.""" + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + with quantized_model_init( + enabled=True, + recipe=recipe.MXFP8BlockScaling(), + preserve_high_precision_init_val=True, + ): + module = GroupedLinear( + num_gemms=2, + in_features=32, + out_features=64, + bias=False, + params_dtype=torch.bfloat16, + single_grouped_weight=False, + ).cuda() + + module.weight0.clear_high_precision_init_val() + with pytest.raises( + RuntimeError, + match="inconsistent high-precision initialization state", + ): + module.make_grouped_weights() + + def test_sanity_checkpointing_on_callables(): """Test that TE checkpointing works correctly on callable modules.""" diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 8ebce563ce..f61e7b4111 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -3,9 +3,11 @@ # See LICENSE for license information. import abc +import contextlib import pytest import torch +from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode try: from torch._opaque_base import OpaqueBaseMeta @@ -27,6 +29,9 @@ from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer from transformer_engine.pytorch.quantization import QuantizerRole +from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer +from transformer_engine.pytorch.quantized_tensor import QuantizedTensor, Quantizer +from transformer_engine.pytorch.dynamo import TensorSpec, to_tensor_spec from transformer_engine.pytorch import ( is_fp8_available, is_mxfp8_available, @@ -37,11 +42,16 @@ NVFP4Quantizer, ) from utils import recipe_id +from transformer_engine.pytorch.attention.dot_product_attention.backends import ( + UnfusedDotProductAttention, +) fp8_available, reason_for_no_fp8 = is_fp8_available(return_reason=True) mxfp8_available, reason_for_no_mxfp8 = is_mxfp8_available(return_reason=True) -fp8_block_scaling_available = is_fp8_block_scaling_available() -nvfp4_available = is_nvfp4_available() +fp8_block_scaling_available, reason_for_no_fp8_block_scaling = is_fp8_block_scaling_available( + return_reason=True +) +nvfp4_available, reason_for_no_nvfp4 = is_nvfp4_available(return_reason=True) def nvfp4_row_scaled(): @@ -389,6 +399,269 @@ def fn(inp): out.sum().backward() +_UNFUSED_DPA_CONFIG = dict( + batch_size=2, + num_heads=4, + head_dim=64, + max_seqlen_q=128, + max_seqlen_kv=128, +) + + +def _make_unfused_attention(dtype: torch.dtype) -> UnfusedDotProductAttention: + cfg = _UNFUSED_DPA_CONFIG + softmax_scale = cfg["head_dim"] ** -0.5 + module = UnfusedDotProductAttention( + softmax_scale=softmax_scale, + attention_type="self", + attention_dropout=0.0, + layer_number=1, + softmax_type="vanilla", + return_max_logit=False, + ) + return module.to(dtype=dtype, device="cuda") + + +_EMPTY_ALIBI_CACHE = { + "_num_heads": None, + "_alibi_slopes": None, + "_max_seqlen_q": None, + "_max_seqlen_kv": None, + "_bottom_right_alignment": True, + "_alibi_bias": None, + "_alibi_slopes_require_update": False, + "_alibi_bias_require_update": False, +} + + +def _make_unfused_qkv(qkv_layout: str, dtype: torch.dtype, requires_grad: bool = True): + """Build (q, k, v) tensors matching `qkv_layout`. Returns also the + extra kwargs (`cu_seqlens_*`, `max_seqlen_*`) that the unfused module + needs for `thd` layouts (empty dict otherwise).""" + cfg = _UNFUSED_DPA_CONFIG + b, s_q, s_kv = cfg["batch_size"], cfg["max_seqlen_q"], cfg["max_seqlen_kv"] + h, d = cfg["num_heads"], cfg["head_dim"] + qkv_format = "".join(c for c in qkv_layout.split("_")[0] if c.isalpha()) + + extra: dict = {} + + def _separate(shape): + return tuple( + torch.randn(shape, dtype=dtype, device="cuda", requires_grad=requires_grad) + for _ in range(3) + ) + + if qkv_layout == "bshd_bshd_bshd": + q, k, v = _separate((b, s_q, h, d)) + elif qkv_layout == "sbhd_sbhd_sbhd": + q, k, v = _separate((s_q, b, h, d)) + elif qkv_layout == "thd_thd_thd": + # All sequences in the batch have the maximum length; no padding. + cu = torch.arange(0, (b + 1) * s_q, step=s_q, dtype=torch.int32, device="cuda") + q, k, v = _separate((b * s_q, h, d)) + extra = dict( + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + max_seqlen_q=s_q, + max_seqlen_kv=s_kv, + ) + elif qkv_layout == "bs3hd": + # Packed: shape (b, s, 3, h, d), q/k/v are views along dim=-3. + qkv = torch.randn( + (b, s_q, 3, h, d), + dtype=dtype, + device="cuda", + requires_grad=requires_grad, + ) + q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + # q/k/v are non-leaf views; retain their grads so the assertions in + # the test (`q.grad is not None` etc.) work for packed layouts. + if requires_grad: + for t in (q, k, v): + t.retain_grad() + elif qkv_layout == "sbh3d": + # Packed: shape (s, b, h, 3, d), q/k/v are views along dim=-2. + qkv = torch.randn( + (s_q, b, h, 3, d), + dtype=dtype, + device="cuda", + requires_grad=requires_grad, + ) + q, k, v = qkv[:, :, :, 0], qkv[:, :, :, 1], qkv[:, :, :, 2] + if requires_grad: + for t in (q, k, v): + t.retain_grad() + else: + raise ValueError(f"Unsupported qkv_layout in test: {qkv_layout}") + + return q, k, v, extra, qkv_format + + +def _call_unfused( + module: UnfusedDotProductAttention, + qkv_layout: str, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + extra: dict, +) -> torch.Tensor: + return module( + _EMPTY_ALIBI_CACHE, + q, + k, + v, + qkv_layout=qkv_layout, + attn_mask_type="causal", + **extra, + ) + + +@pytest.mark.parametrize( + "qkv_layout", + [ + "bshd_bshd_bshd", + "sbhd_sbhd_sbhd", + "thd_thd_thd", + "bs3hd", + "sbh3d", + ], +) +def test_unfused_dpa_torch_compile(qkv_layout): + """Compile UnfusedDotProductAttention.forward with + `torch.compile(fullgraph=True, mode="reduce-overhead")` for several + qkv layouts. + + - `fullgraph=True` makes the test fail on any graph break inside the + unfused attention path. + - `mode="reduce-overhead"` uses the inductor cudagraphs backend, so + forward+backward are captured into CUDA graphs and replayed on + subsequent iterations.""" + dtype = torch.bfloat16 + + module = _make_unfused_attention(dtype) + + def fn(q, k, v, extra): + return _call_unfused(module, qkv_layout, q, k, v, extra) + + torch._dynamo.reset() + compiled = torch.compile(fn, fullgraph=True, mode="reduce-overhead") + + for _ in range(3): + q, k, v, extra, _ = _make_unfused_qkv(qkv_layout, dtype, requires_grad=True) + out = compiled(q, k, v, extra) + out.sum().backward() + torch.cuda.synchronize() + assert torch.isfinite(out).all() + assert q.grad is not None + assert k.grad is not None + assert v.grad is not None + + +# --------------------------------------------------------------------------- +# get_attention_backend under torch.compile +# --------------------------------------------------------------------------- + + +# Scalars in AttentionParams must stay concrete: assume_constant_result cannot +# convert symbolic scalars (dynamo's automatic dynamic would make changed ints +# symbolic on recompilation), so pin them static explicitly. +@torch._dynamo.config.patch(specialize_int=True, specialize_float=True, recompile_limit=32) +def test_get_attention_backend_traceable(monkeypatch): + """get_attention_backend must trace under torch.compile(fullgraph=True) + without graph breaks. The compiled selection must stay consistent with + eager when NVTE_* env vars flip (dynamo guards on os.environ) and when + attention params change, and the baked tex.get_fused_attn_backend result + must drive the selection.""" + from transformer_engine.pytorch.attention.dot_product_attention import utils as dpa_utils + + def fn(x, params): + ( + use_flash_attention, + _, + use_fused_attention, + fused_attention_backend, + use_unfused_attention, + _, + ) = dpa_utils.get_attention_backend(params) + # Encode the full selection (enabled backends + fused sub-backend) in + # the tensor value: without a tensor op dynamo skips the frame entirely + # (nothing gets compiled or guarded), and the output makes compiled vs + # eager selection directly comparable. + return ( + x + + (1 if use_flash_attention else 0) + + (2 if use_fused_attention else 0) + + (4 if use_unfused_attention else 0) + + (8 * int(fused_attention_backend) if fused_attention_backend is not None else 0) + ) + + # Dynamo only guards os.environ entries that exist at trace time (reads of + # absent keys are not guarded yet), so set the vars explicitly. + for env_var, value in ( + ("NVTE_FLASH_ATTN", "1"), + ("NVTE_FUSED_ATTN", "1"), + ("NVTE_UNFUSED_ATTN", "1"), + ("NVTE_FP8_DPA_BWD", "1"), + ("NVTE_DPA_FP8CS_O_in_F16", "1"), + ("NVTE_DPA_FP8_RECIPE", ""), + ("NVTE_UnfusedDPA_Emulate_FP8", "0"), + ): + monkeypatch.setenv(env_var, value) + + torch._dynamo.reset() + compiled = torch.compile(fn, fullgraph=True) + x = torch.zeros(8, device="cuda") + params = dpa_utils.AttentionParams() + + torch.testing.assert_close(compiled(x, params), fn(x, params)) + + # Flip env vars one by one: the compiled function must recompile (guards + # on os.environ) and keep matching eager. + for env_var in ("NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN", "NVTE_FLASH_ATTN"): + monkeypatch.setenv(env_var, "0") + torch.testing.assert_close(compiled(x, params), fn(x, params)) + monkeypatch.setenv(env_var, "1") + + # FP8 attention (fp8_dpa recipe): covers the FP8-only branch (run_config + # env reads, recipe filters, get_fp8_te_dtype). Flipping an FP8-only env + # var (emulation enables UnfusedDotProductAttention) must recompile too. + fp8_params = dpa_utils.AttentionParams( + fp8=True, fp8_meta={"recipe": recipe.DelayedScaling(fp8_dpa=True)} + ) + torch.testing.assert_close(compiled(x, fp8_params), fn(x, fp8_params)) + monkeypatch.setenv("NVTE_UnfusedDPA_Emulate_FP8", "1") + torch.testing.assert_close(compiled(x, fp8_params), fn(x, fp8_params)) + monkeypatch.setenv("NVTE_UnfusedDPA_Emulate_FP8", "0") + + # Changing attention params (ints, layout string, dtype) must recompile + # and keep matching eager, still with no graph break. + for changed_params in ( + dpa_utils.AttentionParams(head_dim_qk=128, head_dim_v=128), + dpa_utils.AttentionParams(max_seqlen_q=512, max_seqlen_kv=512), + dpa_utils.AttentionParams(qkv_layout="bshd_bshd_bshd"), + dpa_utils.AttentionParams(qkv_dtype=torch.float16), + ): + torch.testing.assert_close(compiled(x, changed_params), fn(x, changed_params)) + + # The baked probe result must drive the selection: report no fused + # sub-backend and expect UnfusedDotProductAttention (flash disabled, so the + # outcome is deterministic). Use a fresh frame: already-compiled frames + # keep the previously baked constant (assume_constant_result installs no + # guard on the wrapped function). + monkeypatch.setenv("NVTE_FLASH_ATTN", "0") + monkeypatch.setattr( + dpa_utils.tex, + "get_fused_attn_backend", + lambda *args: dpa_utils.FusedAttnBackend["No_Backend"], + ) + + def fn_no_backend(x, params): + return fn(x, params) + + compiled_no_backend = torch.compile(fn_no_backend, fullgraph=True) + torch.testing.assert_close(compiled_no_backend(x, params), x + 4.0) + + # --------------------------------------------------------------------------- # Value-opaque quantizers # --------------------------------------------------------------------------- @@ -445,14 +718,7 @@ def _hw_available(quantizer): pytest.param(_mxfp8, id="mxfp8"), pytest.param(_blockwise, id="float8_blockwise"), pytest.param(_current_scaling, id="float8_current_scaling"), - pytest.param( - _nvfp4, - id="nvfp4", - marks=pytest.mark.skipif( - not torch.cuda.is_available(), - reason="NVFP4Quantizer requires CUDA to construct", - ), - ), + pytest.param(_nvfp4, id="nvfp4"), ] @@ -466,13 +732,15 @@ def test_quantizer_value_object(factory): rebuilt = eval(repr_str, dict(globals_)) # pylint: disable=eval-used assert rebuilt == a and rebuilt is not a assert hash(rebuilt) == hash(a) + # The deprecated amax-reduction group is never part of the value. + assert getattr(rebuilt, "amax_reduction_group", None) is None # The rebuilt quantizer must also *behave* identically, not just compare # equal: equality only looks at the value key, so a field the kernel needs # but that is absent from the key (e.g. NVFP4's derived ``rht_matrix``) would # slip through the checks above and only blow up at quantize time. Run the # real quantize kernel on both and require bit-exact results. - if torch.cuda.is_available() and _hw_available(a): + if _hw_available(a): x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") torch.testing.assert_close(rebuilt(x).dequantize(), a(x).dequantize(), rtol=0.0, atol=0.0) @@ -541,7 +809,7 @@ def test_quantizer_value_object_fullgraph(factory): unlike merely passing the quantizer through. """ q = factory() - if not (torch.cuda.is_available() and _hw_available(q)): + if not _hw_available(q): pytest.skip("format not supported on this HW") op = _QDQ_OPS[type(q)] @@ -554,3 +822,345 @@ def fn(inp): torch._dynamo.reset() out = torch.compile(fn, fullgraph=True)(x) torch.testing.assert_close(out, ref, rtol=0.0, atol=0.0) + + +# --------------------------------------------------------------------------- +# torch.compile-traceable allocation primitives + TensorSpec +# --------------------------------------------------------------------------- + + +# (factory, logical shape) -- shapes respect MXFP8 (mult. of 32) / blockwise (128) +# / NVFP4 (mult. of 16) constraints. +# Format support is gated at runtime, in the tests that run a kernel; the rest is +# pure Python and works on any HW. +_SPEC_QUANTIZERS = [ + pytest.param(_current_scaling, (4, 8), id="fp8_current_scaling"), + pytest.param(_mxfp8, (64, 128), id="mxfp8"), + pytest.param(_blockwise, (128, 256), id="fp8_blockwise"), + pytest.param(_nvfp4, (64, 128), id="nvfp4"), +] + + +def _build_from_primitives(quantizer, shape, dtype, device="cpu"): + """Assemble a quantized tensor straight from the quantizer primitives: + ``alloc_tensors`` (inner tensors) + ``create_metadata`` (ctx) + the storage's + ``__tensor_unflatten__`` -- i.e. exactly what ``TensorSpec.create_tensor`` + does, but without going through :class:`TensorSpec`. + """ + names = tuple(quantizer.inner_tensor_specs(shape)) + ctx = quantizer.create_metadata(shape, dtype=dtype) + allocated = quantizer.alloc_tensors(shape, device=device) + inner_tensors = {name: allocated[name] for name in names} + storage_cls = ctx["cls"] + # Row-major (contiguous) outer stride for ``__tensor_unflatten__``; ``meta`` + # device computes it without allocating storage. + outer_stride = torch.empty(tuple(shape), device="meta").stride() + return storage_cls.__tensor_unflatten__(inner_tensors, ctx, tuple(shape), outer_stride) + + +def _signature(tensor, names): + """Comparable shape/dtype fingerprint of a tensor and its inner tensors.""" + sig = {"__shape__": tuple(tensor.shape), "__dtype__": tensor.dtype} + for name in names: + inner = getattr(tensor, name) + sig[name] = (tuple(inner.shape), inner.dtype) + return sig + + +def _skip_if_dequantize_unsupported(q): + """Skip when this HW can't run ``dequantize()`` for the quantizer's format. + + ``dequantize()`` runs the real kernel on CUDA, so each format has its own + availability gate (mirrors the ``is_*_available`` checks in test_numerics). + """ + if isinstance(q, MXFP8Quantizer): + if not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + elif isinstance(q, NVFP4Quantizer): + if not nvfp4_available: + pytest.skip(reason_for_no_nvfp4) + elif isinstance(q, Float8BlockQuantizer): + if not fp8_block_scaling_available: + pytest.skip(reason_for_no_fp8_block_scaling) + elif not fp8_available: # Float8 current scaling + pytest.skip(reason_for_no_fp8) + + +# ----- Quantizer primitives ----- + + +@pytest.mark.parametrize("factory, shape", _SPEC_QUANTIZERS) +def test_alloc_tensors_fake(factory, shape): + """``alloc_tensors`` produces FakeTensors with the described shapes/dtypes.""" + q = factory() + specs = q.inner_tensor_specs(shape) + with FakeTensorMode(): + alloc = q.alloc_tensors(shape, device="cpu") + assert set(alloc) == set(specs) + for name, (spec_shape, spec_dtype) in specs.items(): + assert isinstance(alloc[name], FakeTensor) + assert tuple(alloc[name].shape) == tuple(spec_shape) + assert alloc[name].dtype == spec_dtype + + +@pytest.mark.parametrize("factory, shape", _SPEC_QUANTIZERS) +def test_storage_flatten_unflatten_roundtrip(factory, shape): + """Storage ``__tensor_flatten__`` / ``__tensor_unflatten__`` round-trips. + + Build a tensor from ``alloc_tensors`` + ``create_metadata``, flatten it, then + unflatten and verify shape/dtype and every inner buffer match before vs after. + """ + q = factory() + _skip_if_dequantize_unsupported(q) + + tensor = _build_from_primitives(q, shape, torch.bfloat16) + names = tuple(q.inner_tensor_specs(shape)) + # Fill inner tensors with deterministic data (empty() may contain NaNs) so + # the round-trip can be checked by value via dequantize(). + for name in names: + inner = getattr(tensor, name) + inner.copy_(torch.arange(inner.numel(), device=inner.device).reshape(inner.shape)) + before = _signature(tensor, names) + expected = tensor.dequantize() + + flat_names, flat_ctx = tensor.__tensor_flatten__() + assert set(flat_names) == set(names) + inner = {name: getattr(tensor, name) for name in flat_names} + rebuilt = type(tensor).__tensor_unflatten__( + inner, flat_ctx, tuple(tensor.shape), tensor.stride() + ) + + assert isinstance(rebuilt, QuantizedTensor) + assert _signature(rebuilt, flat_names) == before + # The reconstructed tensor dequantizes to the same values. + torch.testing.assert_close(rebuilt.dequantize(), expected, atol=0, rtol=0, equal_nan=True) + + +_USAGE_COMBOS = [ + pytest.param(True, True, id="rowwise_columnwise"), + pytest.param(True, False, id="rowwise_only"), + pytest.param(False, True, id="columnwise_only"), +] + + +@pytest.mark.parametrize("factory, shape", _SPEC_QUANTIZERS) +@pytest.mark.parametrize("rowwise, columnwise", _USAGE_COMBOS) +@pytest.mark.parametrize("internal", [False, True], ids=["wrapper", "internal"]) +def test_python_alloc_matches_cpp_make_empty(factory, shape, rowwise, columnwise, internal): + """Pure-Python allocation is interchangeable with the C++ allocation. + + Builds the same quantized tensor twice: via ``Quantizer.make_empty`` + (``tex.create_empty_quantized_tensor``, the C++ path) and via the Python + primitives ``inner_tensor_specs`` + ``create_metadata`` + ``alloc_tensors`` + + ``__tensor_unflatten__`` (exactly what ``TensorSpec.create_tensor`` + does). Checks: + + * structural parity -- same concrete class, buffer set, per-buffer + shape/dtype/device, logical shape/dtype and flatten context; + * functional parity -- the real C++ quantize kernel writes bit-identical + results into the Python-allocated inner tensors as into the C++-allocated + ones, proving the Python buffer description matches the layout + (padding/alignment) the kernels expect. + """ + + # Two independent, identically-configured quantizers so no state can leak + # between the two allocation paths. + def make_quantizer(): + q = factory() + q.set_usage(rowwise=rowwise, columnwise=columnwise) + q.internal = internal + return q + + q_ref = make_quantizer() + if not _hw_available(q_ref): + pytest.skip("format not supported on this HW") + q_py = make_quantizer() + + ref = q_ref.make_empty(shape, dtype=torch.bfloat16, device="cuda") + py = _build_from_primitives(q_py, shape, torch.bfloat16, device="cuda") + + # --- Structural parity --- + assert type(py) is type(ref) + ref_names, ref_ctx = ref.__tensor_flatten__() + py_names, py_ctx = py.__tensor_flatten__() + assert set(py_names) == set(ref_names) + for name in ref_names: + ref_inner, py_inner = getattr(ref, name), getattr(py, name) + assert tuple(py_inner.shape) == tuple(ref_inner.shape), name + assert py_inner.dtype == ref_inner.dtype, name + assert py_inner.device == ref_inner.device, name + + # Logical shape / dtype (bare storages are not torch.Tensors: they expose + # size() and _dtype instead of .shape / .dtype). + if isinstance(ref, QuantizedTensor): + assert tuple(py.shape) == tuple(ref.shape) == tuple(shape) + assert py.dtype == ref.dtype == torch.bfloat16 + else: + assert tuple(py.size()) == tuple(ref.size()) == tuple(shape) + # pylint: disable=protected-access + assert py._dtype == ref._dtype == torch.bfloat16 + + # Flatten context. The quantizer entry needs special handling: production + # quantizers get a value-based __eq__ from register_value_opaque_quantizer, + # but fall back to field-wise comparison for classes that don't define one + # (plain object.__eq__ is identity, which would spuriously fail). + assert set(py_ctx) == set(ref_ctx) + for key in ("cls", "is_tensor", "requires_grad"): + assert py_ctx[key] == ref_ctx[key], key + ref_kwargs, py_kwargs = ref_ctx["nontensor_kwargs"], py_ctx["nontensor_kwargs"] + assert set(py_kwargs) == set(ref_kwargs) + for key in ref_kwargs: + rv, pv = ref_kwargs[key], py_kwargs[key] + if isinstance(rv, Quantizer) or isinstance(pv, Quantizer): + assert type(pv) is type(rv), key + assert (pv.rowwise_usage, pv.columnwise_usage, pv.internal) == ( + rv.rowwise_usage, + rv.columnwise_usage, + rv.internal, + ), key + if type(rv).__eq__ is not object.__eq__: + assert pv == rv, key + else: + assert pv == rv, key + + # --- Functional parity: run the real C++ quantize kernel into both --- + x = torch.randn(*shape, dtype=torch.bfloat16, device="cuda") + + def _quantize_into(quantizer, dst): + if internal: + # update_quantized() only accepts the wrapper classes; internal + # (bare storage) tensors are filled through the same underlying + # kernel binding directly. + tex.quantize(x, quantizer, dst, None) + else: + quantizer.update_quantized(x, dst) + + # Scale-inv padding is never written by the kernel and both paths allocate it + # uninitialized; zero it so the comparison below covers only kernel output. + for name in ref_names: + getattr(ref, name).zero_() + getattr(py, name).zero_() + + # Some combos are rejected by the quantize kernel itself regardless of who + # allocated the tensor (e.g. FP8 current-scaling columnwise-only on + # TN-capable archs: there is no rowwise data buffer and + # nvte_compute_scale_from_amax asserts on it). Parity then means the + # Python-allocated tensor is rejected the same way -- not a silent skip. + try: + _quantize_into(q_ref, ref) + except RuntimeError: + with pytest.raises(RuntimeError): + _quantize_into(q_py, py) + return + _quantize_into(q_py, py) + for name in ref_names: + torch.testing.assert_close( + getattr(py, name), getattr(ref, name), rtol=0.0, atol=0.0, equal_nan=True + ) + + # Value check through dequantize(). Some layouts cannot dequantize at all + # (e.g. FP8 columnwise-only raises NotImplementedError) -- the C++-allocated + # reference defines what is supported, and when it raises, the bitwise + # buffer equality above already proves value parity. + try: + expected = ref.dequantize() + except NotImplementedError: + expected = None + if expected is not None: + torch.testing.assert_close(py.dequantize(), expected, rtol=0.0, atol=0.0) + + +# ----- TensorSpec ----- + + +@pytest.mark.parametrize("factory, shape", _SPEC_QUANTIZERS) +def test_tensor_spec_matches_primitives(factory, shape): + """TensorSpec is a thin wrapper: its ``create_metadata`` / + ``create_inner_tensors`` / ``create_tensor`` match building everything + directly from the quantizer primitives.""" + q = factory() + spec = TensorSpec(shape=shape, dtype=torch.bfloat16, quantizer=q, device=torch.device("cpu")) + assert spec.is_quantized + + # Metadata matches the quantizer's. + assert spec.create_metadata() == q.create_metadata(shape, dtype=torch.bfloat16) + + # inner_names + create_inner_tensors match inner_tensor_specs. + specs = q.inner_tensor_specs(shape) + names = tuple(specs) + assert spec.inner_names() == names + inner_tensors = spec.create_inner_tensors() + assert len(inner_tensors) == len(names) + for name, inner in zip(names, inner_tensors): + exp_shape, exp_dtype = specs[name] + assert tuple(inner.shape) == tuple(exp_shape) + assert inner.dtype == exp_dtype + + # The assembled tensor matches one built directly from the primitives. + direct = _build_from_primitives(q, shape, torch.bfloat16) + assert _signature(spec.create_tensor(), names) == _signature(direct, names) + + +@pytest.mark.parametrize("factory, shape", _SPEC_QUANTIZERS) +@pytest.mark.parametrize("fake", [False, True], ids=["eager", "fake"]) +def test_tensor_spec_create_tensor(factory, shape, fake): + """``create_tensor`` yields a quantized tensor with the right shape/dtype; + its inner tensors are fake exactly under ``FakeTensorMode``.""" + q = factory() + spec = TensorSpec(shape=shape, dtype=torch.bfloat16, quantizer=q, device=torch.device("cpu")) + with FakeTensorMode() if fake else contextlib.nullcontext(): + out = spec.create_tensor() + assert isinstance(out, QuantizedTensor) + assert tuple(out.shape) == tuple(shape) + assert out.dtype == torch.bfloat16 + for name in spec.inner_names(): + assert isinstance(getattr(out, name), FakeTensor) == fake + + +@pytest.mark.parametrize("factory, shape", _SPEC_QUANTIZERS) +def test_tensor_spec_create_tensor_compiles(factory, shape): + """``TensorSpec.create_tensor`` traces under ``fullgraph=True`` (CPU).""" + q = factory() + + def fn(x): + spec = TensorSpec(shape=tuple(x.shape), dtype=x.dtype, quantizer=q, device=x.device) + t = spec.create_tensor() + acc = x.new_zeros(()) + for name in spec.inner_names(): + acc = acc + getattr(t, name).float().sum() + return acc + + x = torch.zeros(*shape, dtype=torch.bfloat16) + torch._dynamo.reset() + out = torch.compile(fn, fullgraph=True)(x) + assert out.shape == () + + +def test_to_tensor_spec_plain(): + """``to_tensor_spec`` describes a plain tensor.""" + t = torch.empty(2, 3, dtype=torch.float32) + spec = to_tensor_spec(t) + assert not spec.is_quantized + assert spec.shape == (2, 3) + assert spec.dtype == torch.float32 + assert spec.inner_names() == ("data",) + + +@pytest.mark.parametrize("factory, shape", _SPEC_QUANTIZERS) +def test_to_tensor_spec_quantized(factory, shape): + """``to_tensor_spec`` round-trips a quantized tensor back into a spec.""" + q = factory() + tensor = TensorSpec( + shape=shape, dtype=torch.bfloat16, quantizer=q, device=torch.device("cpu") + ).create_tensor() + + spec = to_tensor_spec(tensor) + assert spec.is_quantized + assert spec.shape == tuple(shape) + assert spec.dtype == torch.bfloat16 + # Same buffer layout as the original tensor. + assert spec.inner_names() == tuple(q.inner_tensor_specs(shape)) + # Rebuilding from the derived spec matches the original tensor's structure. + assert _signature(spec.create_tensor(), spec.inner_names()) == _signature( + tensor, spec.inner_names() + ) diff --git a/tests/pytorch/test_weight_swizzle_in_layers.py b/tests/pytorch/test_weight_swizzle_in_layers.py index b9f19f2fd9..ff961c5083 100644 --- a/tests/pytorch/test_weight_swizzle_in_layers.py +++ b/tests/pytorch/test_weight_swizzle_in_layers.py @@ -21,19 +21,25 @@ # skipped individually when the hardware/recipe is unavailable. _SWIZZLING_RECIPES = [ pytest.param( - MXFP8BlockScaling, + "mxfp8", marks=pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8), - id="mxfp8", ), pytest.param( - NVFP4BlockScaling, + "nvfp4-2d", + marks=pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4), + ), + pytest.param( + "nvfp4-1d", marks=pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4), - id="nvfp4", ), ] _LAYER_TYPES = ["Linear", "LayerNormLinear", "LayerNormMLP", "GroupedLinear"] +# 1024 is 128-aligned (NVFP4 swizzle fusion eligible). 1056 is valid MXFP8/NVFP4 +# (divisible by 32) but not 128-aligned, so NVFP4 keeps compact cached scales. +_WEIGHT_SHAPES = [1024, 1056] + def _make_module(layer_type, in_features, out_features, device, num_gemms=1): common = dict(bias=True, params_dtype=torch.bfloat16) @@ -65,12 +71,17 @@ def _grouped_m_splits(layer_type, batch, num_gemms): return [batch // num_gemms] * num_gemms -def _make_recipe(recipe_cls): +def _make_recipe(recipe_name): """Instantiate a recipe with run-to-run nondeterminism disabled where it exists (NVFP4 stochastic rounding); MXFP8 has none.""" - if recipe_cls is NVFP4BlockScaling: - return recipe_cls(disable_stochastic_rounding=True) - return recipe_cls() + if recipe_name == "mxfp8": + return MXFP8BlockScaling() + if recipe_name in ("nvfp4-2d", "nvfp4-1d"): + return NVFP4BlockScaling( + disable_stochastic_rounding=True, + disable_2d_quantization=recipe_name == "nvfp4-1d", + ) + raise ValueError(f"unknown recipe {recipe_name}") def _forward_backward(module, x, is_first_microbatch, recipe, m_splits): @@ -110,21 +121,16 @@ def _run_step(module, x, is_first_microbatch, recipe, m_splits): @pytest.mark.parametrize("layer_type", _LAYER_TYPES) -@pytest.mark.parametrize("recipe_cls", _SWIZZLING_RECIPES) -def test_weight_swizzling_with_workspace_caching(layer_type, recipe_cls): - """Caching the quantized weight across microbatches must pre-swizzle its - scales: the weight quantizer enables ``optimize_for_gemm`` and the cached - workspace carries ``_with_gemm_swizzled_scales``. - - Generic across every swizzling recipe (MXFP8, NVFP4) and module type. - """ +@pytest.mark.parametrize("recipe_name", _SWIZZLING_RECIPES) +def test_weight_swizzling_with_workspace_caching(layer_type, recipe_name): + """Cached weights use preswizzled scales only when direct fusion is supported.""" torch.manual_seed(1234) device = "cuda" in_features, out_features = 1024, 1024 batch = 512 num_gemms = 2 if layer_type == "GroupedLinear" else 1 m_splits = _grouped_m_splits(layer_type, batch, num_gemms) - recipe = recipe_cls() + recipe = _make_recipe(recipe_name) module = _make_module(layer_type, in_features, out_features, device, num_gemms) x = torch.randn(batch, in_features, dtype=torch.bfloat16, device=device, requires_grad=True) @@ -132,11 +138,12 @@ def test_weight_swizzling_with_workspace_caching(layer_type, recipe_cls): # is_first_microbatch=True caches the quantized weight. weight_quantizers = _forward_backward(module, x, True, recipe, m_splits) + expect_preswizzle = recipe_name != "nvfp4-1d" for weight_quantizer in weight_quantizers: assert weight_quantizer is not None assert ( - weight_quantizer.optimize_for_gemm is True - ), f"optimize_for_gemm must be enabled for cached {layer_type} weights" + weight_quantizer.optimize_for_gemm is expect_preswizzle + ), f"unexpected optimize_for_gemm for cached {layer_type} weights" workspaces = module._fp8_workspaces assert len(workspaces) == _expected_num_weights( @@ -144,13 +151,13 @@ def test_weight_swizzling_with_workspace_caching(layer_type, recipe_cls): ), f"unexpected cached weight workspace count for {layer_type}: {len(workspaces)}" for name, ws in workspaces.items(): assert ( - getattr(ws, "_with_gemm_swizzled_scales", False) is True - ), f"cached weight workspace {name!r} scales were not pre-swizzled" + getattr(ws, "_with_gemm_swizzled_scales", False) is expect_preswizzle + ), f"unexpected scale layout for cached weight workspace {name!r}" @pytest.mark.parametrize("layer_type", _LAYER_TYPES) -@pytest.mark.parametrize("recipe_cls", _SWIZZLING_RECIPES) -def test_weight_swizzling_with_primary_fp8_weights(layer_type, recipe_cls): +@pytest.mark.parametrize("recipe_name", _SWIZZLING_RECIPES) +def test_weight_swizzling_with_primary_fp8_weights(layer_type, recipe_name): """With quantized_model_init the weight parameter is itself quantized and is all-gathered (FSDP2) / optimizer-updated in its unswizzled layout, so the eager-swizzle optimization must stay off: the weight quantizer must keep @@ -164,7 +171,7 @@ def test_weight_swizzling_with_primary_fp8_weights(layer_type, recipe_cls): batch = 512 num_gemms = 2 if layer_type == "GroupedLinear" else 1 m_splits = _grouped_m_splits(layer_type, batch, num_gemms) - recipe = recipe_cls() + recipe = _make_recipe(recipe_name) with te.quantized_model_init(enabled=True, recipe=recipe): module = _make_module(layer_type, in_features, out_features, device, num_gemms) @@ -191,8 +198,35 @@ def test_weight_swizzling_with_primary_fp8_weights(layer_type, recipe_cls): @pytest.mark.parametrize("layer_type", _LAYER_TYPES) -@pytest.mark.parametrize("recipe_cls", _SWIZZLING_RECIPES) -def test_weight_caching_matches_no_caching(layer_type, recipe_cls): +@pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4) +def test_weight_optimize_for_gemm_disabled_without_swizzle_fusion(layer_type): + """NVFP4 weights that cannot use in-kernel swizzle fusion must keep compact cached scales.""" + torch.manual_seed(1234) + device = "cuda" + # Valid for NVFP4 quantization, but not aligned to 128x128 swizzle tiles. + in_features, out_features = 1056, 1056 + batch = 512 + num_gemms = 2 if layer_type == "GroupedLinear" else 1 + m_splits = _grouped_m_splits(layer_type, batch, num_gemms) + recipe = NVFP4BlockScaling(disable_stochastic_rounding=True) + + module = _make_module(layer_type, in_features, out_features, device, num_gemms) + x = torch.randn(batch, in_features, dtype=torch.bfloat16, device=device, requires_grad=True) + + weight_quantizers = _forward_backward(module, x, True, recipe, m_splits) + + for weight_quantizer in weight_quantizers: + assert weight_quantizer is not None + assert weight_quantizer.optimize_for_gemm is False + + for _, ws in module._fp8_workspaces.items(): + assert getattr(ws, "_with_gemm_swizzled_scales", False) is False + + +@pytest.mark.parametrize("layer_type", _LAYER_TYPES) +@pytest.mark.parametrize("recipe_name", _SWIZZLING_RECIPES) +@pytest.mark.parametrize("features", _WEIGHT_SHAPES) +def test_weight_caching_matches_no_caching(layer_type, recipe_name, features): """Caching the quantized (pre-swizzled) weight across microbatches must be numerically identical to the uncached flow that re-quantizes the weight every microbatch. Verified per microbatch for fprop output, dgrad and wgrad, across @@ -205,12 +239,12 @@ def test_weight_caching_matches_no_caching(layer_type, recipe_cls): """ torch.manual_seed(1234) device = "cuda" - in_features, out_features = 1024, 1024 + in_features, out_features = features, features batch = 512 microbatches = 4 num_gemms = 2 if layer_type == "GroupedLinear" else 1 m_splits = _grouped_m_splits(layer_type, batch, num_gemms) - recipe = _make_recipe(recipe_cls) + recipe = _make_recipe(recipe_name) # Identical modules: one drives the cached path, one the uncached path. cached = _make_module(layer_type, in_features, out_features, device, num_gemms) diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index b1ef46b96e..5be2d6c07a 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -120,6 +120,7 @@ def quantization_tols(name: str) -> dict[str, float]: "fp8_delayed_scaling", "fp8_current_scaling", "fp8_blockwise", + "fp8_block_scaling", "mxfp8", "mxfp8_block_scaling", ): diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 90d5d5679f..8353c42171 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -36,10 +36,6 @@ if (CMAKE_BUILD_TYPE STREQUAL "Debug") set(CMAKE_CUDA_FLAGS_DEBUG "${CMAKE_CUDA_FLAGS_DEBUG} -g -G") endif() -# Hide non-necessary symbols in shared object. -set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wl,--version-script=${CMAKE_CURRENT_SOURCE_DIR}/libtransformer_engine.version") -set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Wl,--version-script=${CMAKE_CURRENT_SOURCE_DIR}/libtransformer_engine.version") - # Transformer Engine library project(transformer_engine LANGUAGES CUDA CXX) @@ -77,6 +73,9 @@ if(NOT arch_100_index EQUAL -1) if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.9) list(APPEND NVTE_SPECIFIC_ARCHS "103a") endif() + if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 13.4) + list(APPEND NVTE_SPECIFIC_ARCHS "107a") + endif() endif() # Check for architecture 101 (if we see this we are in toolkit <= 12.9) @@ -158,17 +157,20 @@ endif() # Python find_package(Python COMPONENTS Interpreter Development.Module REQUIRED) -function(find_nccl_version OUT_VERSION OUT_INCLUDE_DIR) - find_path(_nvte_nccl_include_dir +find_path(NCCL_INCLUDE_DIR NAMES nccl.h + HINTS "${Python_SITEARCH}/nvidia/nccl" + "/opt/nvidia/nccl" + "/usr/local/nccl" PATH_SUFFIXES include REQUIRED) - file(STRINGS "${_nvte_nccl_include_dir}/nccl.h" _nvte_nccl_major_line +function(get_nccl_version OUT_VERSION INCLUDE_DIR) + file(STRINGS "${INCLUDE_DIR}/nccl.h" _nvte_nccl_major_line REGEX "^#define NCCL_MAJOR[ \t]+[0-9]+$") - file(STRINGS "${_nvte_nccl_include_dir}/nccl.h" _nvte_nccl_minor_line + file(STRINGS "${INCLUDE_DIR}/nccl.h" _nvte_nccl_minor_line REGEX "^#define NCCL_MINOR[ \t]+[0-9]+$") - file(STRINGS "${_nvte_nccl_include_dir}/nccl.h" _nvte_nccl_patch_line + file(STRINGS "${INCLUDE_DIR}/nccl.h" _nvte_nccl_patch_line REGEX "^#define NCCL_PATCH[ \t]+[0-9]+$") string(REGEX REPLACE "^#define NCCL_MAJOR[ \t]+([0-9]+)$" "\\1" @@ -182,15 +184,16 @@ function(find_nccl_version OUT_VERSION OUT_INCLUDE_DIR) OR "${_nvte_nccl_minor}" STREQUAL "" OR "${_nvte_nccl_patch}" STREQUAL "") message(FATAL_ERROR - "Failed to parse NCCL version from ${_nvte_nccl_include_dir}/nccl.h") + "Failed to parse NCCL version from ${INCLUDE_DIR}/nccl.h") endif() set(${OUT_VERSION} "${_nvte_nccl_major}.${_nvte_nccl_minor}.${_nvte_nccl_patch}" PARENT_SCOPE) - set(${OUT_INCLUDE_DIR} "${_nvte_nccl_include_dir}" PARENT_SCOPE) endfunction() +get_nccl_version(NCCL_VERSION "${NCCL_INCLUDE_DIR}") + function(find_cublasmp_version OUT_VERSION OUT_INCLUDE_DIR SEARCH_DIR) find_path(_nvte_cublasmp_include_dir NAMES cublasmp.h @@ -434,6 +437,14 @@ foreach(cuda_source IN LISTS transformer_engine_cuda_arch_specific_sources) endforeach() add_library(transformer_engine SHARED ${transformer_engine_SOURCES}) + +# This is TE-specific and should not apply to all targets +target_link_options( + transformer_engine + PRIVATE + "LINKER:--version-script=${CMAKE_CURRENT_SOURCE_DIR}/libtransformer_engine.version" +) + # Disable CMake's automatic architecture flag injection. # All architectures are handled explicitly via per-source COMPILE_OPTIONS # using NVTE_STANDARD_ARCHS, NVTE_GENERIC_ARCHS, and NVTE_SPECIFIC_ARCHS above. @@ -507,9 +518,11 @@ target_link_libraries(transformer_engine PUBLIC CUDA::cublas CUDA::cudart CUDNN::cudnn_all) +target_link_libraries(transformer_engine PRIVATE ${CMAKE_DL_LIBS}) target_include_directories(transformer_engine PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) +target_include_directories(transformer_engine PRIVATE ${NCCL_INCLUDE_DIR}) target_include_directories(transformer_engine SYSTEM PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}/cccl) target_include_directories(transformer_engine PRIVATE "${CUDNN_FRONTEND_INCLUDE_DIR}") @@ -543,7 +556,6 @@ if (NVTE_WITH_CUBLASMP) target_compile_definitions(transformer_engine PRIVATE NVTE_WITH_CUBLASMP) target_include_directories(transformer_engine PRIVATE ${CUBLASMP_DIR}/include) - find_nccl_version(NCCL_VERSION NCCL_INCLUDE_DIR) find_cublasmp_version(CUBLASMP_VERSION CUBLASMP_INCLUDE_DIR ${CUBLASMP_DIR}) find_library(CUBLASMP_LIB NAMES cublasmp libcublasmp.so libcublasmp.so.0 @@ -598,10 +610,10 @@ option(NVTE_WITH_NCCL_EP "Build NCCL EP into libtransformer_engine.so" ON) if(NVTE_WITH_NCCL_EP) # SM>=90 and NCCL>=2.30.4 are gated at runtime in EPBackend::initialize. # -- NCCL EP headers -------------------------------------------------------- -# Headers + libs are produced by the in-tree 3rdparty/nccl submodule build +# Headers + libs are produced by the in-tree 3rdparty/nccl-extensions submodule build # (auto-built by setup.py via build_nccl_ep_submodule). set(NCCL_EP_SUBMODULE_ROOT - "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl") + "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl-extensions") set(NCCL_EP_INCLUDE_DIR "${NCCL_EP_SUBMODULE_ROOT}/build/include") if(NOT EXISTS "${NCCL_EP_INCLUDE_DIR}/nccl_ep.h") message(FATAL_ERROR @@ -622,15 +634,7 @@ find_file(NCCL_EP_LIB NO_DEFAULT_PATH REQUIRED) -# -- NCCL core: nccl.h + libnccl.so ----------------------------------------- -# setup.py passes -DNCCL_INCLUDE_DIR; standalone CMake falls back to probing -# well-known NCCL install prefixes. -find_path(NCCL_INCLUDE_DIR nccl.h - HINTS /opt/nvidia/nccl/include /usr/local/nccl/include) -if(NOT NCCL_INCLUDE_DIR) - message(FATAL_ERROR - "nccl.h not found. Pass -DNCCL_INCLUDE_DIR=/include.") -endif() +# -- NCCL core library ------------------------------------------------------- if(NOT NCCL_LIB) find_library(NCCL_LIB NAMES nccl libnccl @@ -639,8 +643,7 @@ if(NOT NCCL_LIB) endif() target_include_directories(transformer_engine PRIVATE - ${NCCL_EP_INCLUDE_DIR} - ${NCCL_INCLUDE_DIR}) + ${NCCL_EP_INCLUDE_DIR}) # libnccl.so direct symbols (ncclGetVersion etc.) come from libnccl_ep.a's # DT_NEEDED chain plus this TU's own references. CUDA::cuda_driver must follow diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index 6f8af7d8f0..b2dffa567e 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -282,24 +282,49 @@ def _get_sys_extension() -> str: raise RuntimeError(f"Unsupported operating system ({system})") +def _cuda_runtime_major(cuda_runtime: ctypes.CDLL) -> Optional[int]: + """Return the major version of the CUDA runtime loaded with Transformer Engine.""" + + runtime_version = ctypes.c_int() + get_runtime_version = cuda_runtime.cudaRuntimeGetVersion + get_runtime_version.argtypes = [ctypes.POINTER(ctypes.c_int)] + get_runtime_version.restype = ctypes.c_int + if get_runtime_version(ctypes.byref(runtime_version)) != 0 or runtime_version.value <= 0: + return None + return runtime_version.value // 1000 + + @functools.lru_cache(maxsize=None) -def _nvidia_cudart_include_dir() -> str: +def _nvidia_cudart_include_dir(cuda_major_version: int) -> str: """Returns the include directory for cuda_runtime.h if exists in python environment.""" + # This is primarily here to support editable installs. cuda_runtime.cpp handles the + # resolution for install via wheel or when using the shared library ABI directly. + try: import nvidia except ModuleNotFoundError: return "" - # Installing some nvidia-* packages, like nvshmem, create nvidia name, so "import nvidia" - # above doesn't throw. However, they don't set "__file__" attribute. + # NVIDIA packages may use either a regular package or a namespace package spread + # across multiple package roots. if nvidia.__file__ is not None: - nvidia_root = Path(nvidia.__file__).parent + nvidia_roots = (Path(nvidia.__file__).parent,) else: - nvidia_root = Path(nvidia.__path__[0]) # namespace package + nvidia_roots = tuple(Path(path) for path in nvidia.__path__) + + layouts = [f"cu{cuda_major_version}"] + if cuda_major_version == 12: + layouts.append("cuda_runtime") - include_dir = nvidia_root / "cuda_runtime" - return str(include_dir) if include_dir.exists() else "" + for layout in layouts: + for nvidia_root in nvidia_roots: + cuda_root = nvidia_root / layout + if (cuda_root / "cuda_runtime.h").is_file() or ( + cuda_root / "include" / "cuda_runtime.h" + ).is_file(): + return str(cuda_root) + return "" @functools.lru_cache(maxsize=None) @@ -470,5 +495,8 @@ def _load_core_library(): pass else: # Needed to find the correct headers for NVRTC kernels. - if not os.getenv("NVTE_CUDA_INCLUDE_DIR") and _nvidia_cudart_include_dir(): - os.environ["NVTE_CUDA_INCLUDE_DIR"] = _nvidia_cudart_include_dir() + _cuda_major_version = _cuda_runtime_major(_TE_LIB_CTYPES) + if not os.getenv("NVTE_CUDA_INCLUDE_DIR") and _cuda_major_version is not None: + cuda_include_dir = _nvidia_cudart_include_dir(_cuda_major_version) + if cuda_include_dir: + os.environ["NVTE_CUDA_INCLUDE_DIR"] = cuda_include_dir diff --git a/transformer_engine/common/cast/cast_grouped_dbias.cu b/transformer_engine/common/cast/cast_grouped_dbias.cu index 5290255a00..b7ced30b11 100644 --- a/transformer_engine/common/cast/cast_grouped_dbias.cu +++ b/transformer_engine/common/cast/cast_grouped_dbias.cu @@ -11,7 +11,8 @@ #include "dispatch/quantize.cuh" void nvte_group_quantize_dbias(const NVTEGroupedTensor input, NVTEGroupedTensor output, - NVTEGroupedTensor dbias, NVTETensor workspace, cudaStream_t stream) { + NVTEGroupedTensor dbias, NVTETensor workspace, + const NVTEQuantizationConfig quant_config, cudaStream_t stream) { NVTE_API_CALL(nvte_group_quantize_dbias); using namespace transformer_engine; @@ -20,5 +21,5 @@ void nvte_group_quantize_dbias(const NVTEGroupedTensor input, NVTEGroupedTensor constexpr const NVTEGroupedTensor activation_input = nullptr; dispatch::group_quantize_bwd_helper( - input, activation_input, output, dbias, workspace, nullptr, stream); + input, activation_input, output, dbias, workspace, quant_config, stream); } diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 79975b288e..51bb4c81b2 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -116,7 +116,7 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, Tensor *dummy_workspace_tensor = nullptr; mxfp8::quantize( *input_tensor, dummy_input_tensor, noop_tensor, output_tensor, dummy_dbias_tensor, - dummy_workspace_tensor, stream); + dummy_workspace_tensor, quant_config_cpp.mxfp8_2d_quantization, stream); break; } #ifdef __HIP_PLATFORM_AMD__ @@ -155,9 +155,20 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, if (row_scaled_nvfp4) { NVTE_CHECK(!quant_config_cpp.nvfp4_2d_quantization, "Row-scaled NVFP4 quantization does not support 2D quantization."); - NVTE_CHECK(!output_tensor->has_columnwise_data(), - "Row-scaled NVFP4 quantization does not produce columnwise output."); + NVTE_CHECK( + !(nvfp4_use_4over6 && output_tensor->has_columnwise_data()), + "Row-scaled NVFP4 transpose quantization is not supported with 4over6 mode. The 4over6 " + "kernel does not consume the per-row/per-column amaxes, so the columnwise output would " + "be incorrect."); + NVTE_CHECK( + !output_tensor->has_columnwise_data() || + (dtype == DType::kBFloat16 && rows % 32 == 0 && cols % 32 == 0), + "Row-scaled NVFP4 transpose quantization requires BF16 input and dimensions that are " + "multiples of 32."); nvfp4::compute_rowwise_amax(*input_tensor, noop_tensor, output_tensor, stream); + if (output_tensor->has_columnwise_data()) { + nvfp4::compute_columnwise_amax(*input_tensor, noop_tensor, output_tensor, stream); + } } // Columnwise-only is supported on the optimized path only for 2D scaling; rowwise-only and // both-directions keep their existing routing. Columnwise-only 1D and non-bf16 fall back to @@ -303,7 +314,7 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens case NVTE_MXFP8_1D_SCALING: { mxfp8::quantize( *grad_tensor, input_tensor, noop_tensor, output_tensor, dbias_tensor, workspace_tensor, - stream); + quant_config_cpp.mxfp8_2d_quantization, stream); break; } #ifdef __HIP_PLATFORM_AMD__ @@ -326,6 +337,7 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens // Choose kernel const auto [rows, cols] = grad_tensor->flat_2d_dims(); auto dtype = grad_tensor->dtype(); + const bool row_scaled_nvfp4 = output_tensor->row_scaled_nvfp4; const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; NVTE_CHECK(nvfp4_use_4over6 || output_tensor->nvfp4_e4m3_max == 448, "Non-4over6 NVFP4 quantization requires E4M3 max 448."); @@ -335,9 +347,29 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens // Refuse the fast-math error path rather than silently scoring with the exact one. NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.nvfp4_4over6_err_use_fast_math, "NVFP4 4over6 fast-math error mode is not supported on ROCm."); -#endif + // Row-scaled NVFP4 backward is not supported on ROCm; keep refusing it here even though the + // upstream CUDA path below now handles it. NVTE_CHECK(!output_tensor->row_scaled_nvfp4, "Backward NVFP4 quantization does not support row-scaled outputs."); +#endif + if (row_scaled_nvfp4) { + NVTE_CHECK(!quant_config_cpp.nvfp4_2d_quantization, + "Row-scaled NVFP4 quantization does not support 2D quantization."); + NVTE_CHECK( + !(nvfp4_use_4over6 && output_tensor->has_columnwise_data()), + "Row-scaled NVFP4 transpose quantization is not supported with 4over6 mode. The 4over6 " + "kernel does not consume the per-row/per-column amaxes, so the columnwise output would " + "be incorrect."); + NVTE_CHECK( + !output_tensor->has_columnwise_data() || + (dtype == DType::kBFloat16 && rows % 32 == 0 && cols % 32 == 0), + "Row-scaled NVFP4 transpose quantization requires BF16 input and dimensions that are " + "multiples of 32."); + nvfp4::compute_rowwise_amax(*grad_tensor, noop_tensor, output_tensor, stream); + if (output_tensor->has_columnwise_data()) { + nvfp4::compute_columnwise_amax(*grad_tensor, noop_tensor, output_tensor, stream); + } + } // Columnwise-only is supported on the optimized path only for 2D scaling; rowwise-only and // both-directions keep their existing routing. Columnwise-only 1D and non-bf16 fall back to // quantize_transpose_vector_blockwise_fp4. @@ -379,7 +411,7 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens /*use_stochastic_rounding=*/quant_config_cpp.stochastic_rounding, /*rng_state=*/quant_config_cpp.rng_state, /*use_2d_quantization=*/quant_config_cpp.nvfp4_2d_quantization, - /*row_scaled_nvfp4=*/false, + /*row_scaled_nvfp4=*/row_scaled_nvfp4, /*noop_tensor=*/noop_tensor->data, /*nvfp4_e4m3_max=*/output_tensor->nvfp4_e4m3_max, /*nvfp4_4over6_mode=*/quant_config_cpp.nvfp4_4over6_mode, @@ -551,13 +583,10 @@ void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor } case NVTE_BLOCK_SCALING_1D: { NVTE_CHECK(!IS_ACT, "IS_ACT is not implemented for grouped NVTE_BLOCK_SCALING_1D."); - NVTE_CHECK(!quant_config_cpp.force_pow_2_scales, - "Fused grouped FP8 block-scaling quantize does not support " - "force_pow_2_scales=True. Set force_pow_2_scales=False, or use the unfused " - "split-quantize path (NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=0)."); #ifndef __HIP_PLATFORM_AMD__ fp8_blockwise::group_quantize_blockwise_1d(input_tensor, output_tensor, noop_tensor, - quant_config_cpp.amax_epsilon, stream); + quant_config_cpp.amax_epsilon, + quant_config_cpp.force_pow_2_scales, stream); #else NVTE_ERROR( "Grouped FP8 block-scaling quantization is not supported on ROCm platform."); @@ -566,13 +595,10 @@ void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor } case NVTE_BLOCK_SCALING_2D: { NVTE_CHECK(!IS_ACT, "IS_ACT is not implemented for grouped NVTE_BLOCK_SCALING_2D."); - NVTE_CHECK(!quant_config_cpp.force_pow_2_scales, - "Fused grouped FP8 block-scaling quantize does not support " - "force_pow_2_scales=True. Set force_pow_2_scales=False, or use the unfused " - "split-quantize path (NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=0)."); #ifndef __HIP_PLATFORM_AMD__ fp8_blockwise::group_quantize_blockwise_2d(input_tensor, output_tensor, noop_tensor, - quant_config_cpp.amax_epsilon, stream); + quant_config_cpp.amax_epsilon, + quant_config_cpp.force_pow_2_scales, stream); #else NVTE_ERROR( "Grouped FP8 block-scaling quantization is not supported on ROCm platform."); @@ -623,23 +649,19 @@ void group_quantize_bwd_helper(const NVTEGroupedTensor grad, const NVTEGroupedTe case NVTE_BLOCK_SCALING_1D: case NVTE_BLOCK_SCALING_2D: { NVTE_CHECK(!IS_DACT, "IS_DACT is not implemented for grouped FP8 block scaling."); - NVTE_CHECK(!quant_config_cpp.force_pow_2_scales, - "Fused grouped FP8 block-scaling quantize does not support " - "force_pow_2_scales=True. Set force_pow_2_scales=False, or use the unfused " - "split-quantize path (NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=0)."); #ifndef __HIP_PLATFORM_AMD__ // dbias is computed in-kernel and reduced per-expert inside group_quantize_blockwise_{1d,2d} // (mirrors MXFP8); those also handle the two-call workspace sizing protocol. GroupedTensor *dbias_arg = IS_DBIAS ? dbias_tensor : nullptr; Tensor *workspace_arg = IS_DBIAS ? workspace_tensor : nullptr; if (scaling_mode == NVTE_BLOCK_SCALING_1D) { - fp8_blockwise::group_quantize_blockwise_1d(grad_tensor, output_tensor, noop_tensor, - quant_config_cpp.amax_epsilon, stream, dbias_arg, - workspace_arg); + fp8_blockwise::group_quantize_blockwise_1d( + grad_tensor, output_tensor, noop_tensor, quant_config_cpp.amax_epsilon, + quant_config_cpp.force_pow_2_scales, stream, dbias_arg, workspace_arg); } else { - fp8_blockwise::group_quantize_blockwise_2d(grad_tensor, output_tensor, noop_tensor, - quant_config_cpp.amax_epsilon, stream, dbias_arg, - workspace_arg); + fp8_blockwise::group_quantize_blockwise_2d( + grad_tensor, output_tensor, noop_tensor, quant_config_cpp.amax_epsilon, + quant_config_cpp.force_pow_2_scales, stream, dbias_arg, workspace_arg); } #else NVTE_ERROR( diff --git a/transformer_engine/common/cast/fp8/gated_fp8.cuh b/transformer_engine/common/cast/fp8/gated_fp8.cuh index cec1488984..85efb63ef0 100644 --- a/transformer_engine/common/cast/fp8/gated_fp8.cuh +++ b/transformer_engine/common/cast/fp8/gated_fp8.cuh @@ -73,11 +73,9 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) const float scale = (scale_ptr != nullptr) ? *scale_ptr : 1; extern __shared__ char dynamic_shmem[]; - uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); // Manually align dynamic SHMEM per TMA requirements using padding // __align__(128) Does not guarantee the pointer to be aligned! - uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) & - ~(static_cast(TMA_SHMEM_ALIGNMENT - 1)); + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); constexpr size_t buff_elems = SHMEM_DIM_Y * SHMEM_DIM_X; constexpr size_t buff_elems_total = BUFFERS_NUM * buff_elems; diff --git a/transformer_engine/common/cast/fp8/quantize_fp8.cuh b/transformer_engine/common/cast/fp8/quantize_fp8.cuh index 672f1e88f4..b5d13b5b36 100644 --- a/transformer_engine/common/cast/fp8/quantize_fp8.cuh +++ b/transformer_engine/common/cast/fp8/quantize_fp8.cuh @@ -307,9 +307,14 @@ template __global__ void __launch_bounds__(THREADS_PER_BLOCK) cast_fp8_1D_kernel(const IType *input_ptr, OType *output_ptr, float *const amax_ptr, - float *const scale_inv_ptr, const float *const scale_ptr, const size_t N) { + float *const scale_inv_ptr, const float *const scale_ptr, const size_t N, + const float *const noop) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + const size_t block_offset = blockIdx.x * ELEMS_PER_BLOCK; const IType *input = input_ptr + block_offset; OType *output = output_ptr + block_offset; @@ -409,7 +414,7 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK) } // namespace quantize_1D_kernel template -void quantize_1D(const Tensor &input, Tensor *output, cudaStream_t stream) { +void quantize_1D(const Tensor &input, const Tensor *noop, Tensor *output, cudaStream_t stream) { using namespace quantize_1D_kernel; const size_t N = product(input.data.shape); @@ -424,6 +429,7 @@ void quantize_1D(const Tensor &input, Tensor *output, cudaStream_t stream) { float *const amax_ptr = reinterpret_cast(output->amax.dptr); float *const scale_inv_ptr = reinterpret_cast(output->scale_inv.dptr); const float *const scale_ptr = reinterpret_cast(output->scale.dptr); + const float *noop_ptr = reinterpret_cast(noop->data.dptr); const dim3 block(THREADS_PER_BLOCK); const dim3 grid(blocks); @@ -435,9 +441,10 @@ void quantize_1D(const Tensor &input, Tensor *output, cudaStream_t stream) { const IType *input_ptr = reinterpret_cast(input.data.dptr); OType *output_ptr = reinterpret_cast(output->data.dptr); - cast_fp8_1D_kernel<<>>( - input_ptr, output_ptr, amax_ptr, scale_inv_ptr, scale_ptr, N);); // NOLINT(*) - ); // NOLINT(*) + cast_fp8_1D_kernel + <<>>(input_ptr, output_ptr, amax_ptr, scale_inv_ptr, scale_ptr, N, + noop_ptr);); // NOLINT(*) + ); // NOLINT(*) NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -695,7 +702,7 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, is_aligned_tensor_data(input, TMA_GMEM_ALIGNMENT) && is_aligned_tensor_data(*output, TMA_GMEM_ALIGNMENT)) { // Aligned AND FP8 - quantize_1D(input, output, stream); + quantize_1D(input, noop, output, stream); } else { // Unaligned CastVectorizedUnaryKernelLauncher(input, noop, output, stream); diff --git a/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh index 1fd1738f93..203f569471 100644 --- a/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh +++ b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh @@ -284,9 +284,12 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 4) group_block_scaled_2d_tma CType* __restrict__ scale_inv_t_base, const int64_t* __restrict__ tensor_offsets_ptr, const size_t num_tensors, const size_t common_first_dim_blocks, const size_t K, const size_t total_row_blocks, const size_t blocks_X, const size_t scale_stride_y, - const float epsilon, const float* __restrict__ noop_ptr, float* __restrict__ dbias_workspace) { + const float epsilon, const bool pow_2_scales, const float* __restrict__ noop_ptr, + float* __restrict__ dbias_workspace) { #if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 - if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) return; + // Skipping is only safe without dbias: the grouped_reduce_dbias launch is unconditional, so an + // early return would leave it reducing a workspace this kernel never wrote. + if (dbias_workspace == nullptr && noop_ptr != nullptr && noop_ptr[0] == 1.0f) return; const size_t tile_x = blockIdx.x; const size_t tile_y_global = blockIdx.y; @@ -378,8 +381,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 4) group_block_scaled_2d_tma for (int w = 1; w < kNumWarps; ++w) { block_amax = fmaxf(block_amax, warp_amaxes[w]); } - const CType scale = - compute_scale_from_types(block_amax, epsilon, /*pow_2_scaling=*/false); + const CType scale = compute_scale_from_types(block_amax, epsilon, pow_2_scales); // The 2D colwise per-expert scale offset requires a CTA-cooperative prefix // sum in the VARYING_FIRST_DIM case, so compute it across all threads before @@ -475,7 +477,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) const size_t num_tensors, const size_t common_first_dim_blocks, const size_t K, const size_t total_row_blocks, const size_t R_total, const float epsilon, - const float* __restrict__ noop_ptr) { + const bool pow_2_scales, const float* __restrict__ noop_ptr) { #if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) return; @@ -530,8 +532,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) CType amax = compute_row_amax(in_vec[it]); amax = subwarp_reduce_max_broadcast(amax); - const CType scale = - compute_scale_from_types(amax, epsilon, /*pow_2_scaling=*/false); + const CType scale = compute_scale_from_types(amax, epsilon, pow_2_scales); const CType scale_inv = 1.f / scale; if (thr_col == 0 && r_global < R_total) { // Per-expert layout: (blocks_X, roundup(M_t, 4)). Compute expert base @@ -572,10 +573,12 @@ __global__ void __launch_bounds__(kThreadsPerBlock) group_block_scaled_1d_tma_ke CType* __restrict__ scale_inv_t_base, const int64_t* __restrict__ tensor_offsets_ptr, const size_t num_tensors, const size_t common_first_dim_blocks, const size_t K, const size_t total_row_blocks, const size_t blocks_X, const size_t scale_t_stride_aligned_K, - const size_t R_total, const float epsilon, const float* __restrict__ noop_ptr, - float* __restrict__ dbias_workspace) { + const size_t R_total, const float epsilon, const bool pow_2_scales, + const float* __restrict__ noop_ptr, float* __restrict__ dbias_workspace) { #if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000 - if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) return; + // Skipping is only safe without dbias: the grouped_reduce_dbias launch is unconditional, so an + // early return would leave it reducing a workspace this kernel never wrote. + if (dbias_workspace == nullptr && noop_ptr != nullptr && noop_ptr[0] == 1.0f) return; const size_t tile_x = blockIdx.x; const size_t tile_y_global = blockIdx.y; @@ -656,8 +659,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) group_block_scaled_1d_tma_ke CType amax = compute_row_amax(in_vec); amax = subwarp_reduce_max_broadcast(amax); - const CType scale = - compute_scale_from_types(amax, epsilon, /*pow_2_scaling=*/false); + const CType scale = compute_scale_from_types(amax, epsilon, pow_2_scales); const CType scale_inv = 1.f / scale; const size_t r_global = global_row_base + row_local; @@ -733,8 +735,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) group_block_scaled_1d_tma_ke } amax = subwarp_reduce_max_broadcast(amax); - const CType scale = - compute_scale_from_types(amax, epsilon, /*pow_2_scaling=*/false); + const CType scale = compute_scale_from_types(amax, epsilon, pow_2_scales); const CType scale_inv = 1.f / scale; const size_t c_global = global_col_base + col_local; @@ -814,7 +815,8 @@ inline GroupedBlockwiseLaunchInfo prepare_grouped_blockwise_launch(const Grouped // reports the [total_row_blocks, K] fp32 shape and returns without launching. inline void group_quantize_blockwise_2d(const GroupedTensor* input, GroupedTensor* output, const Tensor* noop, const float epsilon, - cudaStream_t stream, GroupedTensor* dbias = nullptr, + const bool pow_2_scales, cudaStream_t stream, + GroupedTensor* dbias = nullptr, Tensor* workspace = nullptr) { const int sm = transformer_engine::cuda::sm_arch(); NVTE_CHECK(sm >= 90 && sm < 100, @@ -883,7 +885,7 @@ inline void group_quantize_blockwise_2d(const GroupedTensor* input, GroupedTenso : nullptr, info.tensor_offsets_d, info.num_tensors, info.common_first_dim_blocks, info.K, info.total_row_blocks, info.blocks_X, scale_stride_y, epsilon, - noop_ptr, dbias_workspace); + pow_2_scales, noop_ptr, dbias_workspace); if (dbias_workspace != nullptr) { const ShapeRepresentation shape_rep = info.same_both_dims ? ShapeRepresentation::SAME_BOTH_DIMS @@ -907,7 +909,8 @@ inline void group_quantize_blockwise_2d(const GroupedTensor* input, GroupedTenso // per-tile column partial can be computed. inline void group_quantize_blockwise_1d(const GroupedTensor* input, GroupedTensor* output, const Tensor* noop, const float epsilon, - cudaStream_t stream, GroupedTensor* dbias = nullptr, + const bool pow_2_scales, cudaStream_t stream, + GroupedTensor* dbias = nullptr, Tensor* workspace = nullptr) { const int sm = transformer_engine::cuda::sm_arch(); NVTE_CHECK(sm >= 90 && sm < 100, @@ -967,7 +970,7 @@ inline void group_quantize_blockwise_1d(const GroupedTensor* input, GroupedTenso reinterpret_cast(output->scale_inv.dptr), info.tensor_offsets_d, info.num_tensors, info.common_first_dim_blocks, info.K, info.total_row_blocks, - info.R_total, epsilon, noop_ptr); + info.R_total, epsilon, pow_2_scales, noop_ptr); } } else if constexpr (kRowwise || kColwise) { // CW-only, BOTH, or RW-only WITH dbias: smem-cached TMA kernel. @@ -998,7 +1001,7 @@ inline void group_quantize_blockwise_1d(const GroupedTensor* input, GroupedTenso : nullptr, info.tensor_offsets_d, info.num_tensors, info.common_first_dim_blocks, info.K, info.total_row_blocks, info.blocks_X, scale_t_stride_aligned_K, - info.R_total, epsilon, noop_ptr, dbias_workspace); + info.R_total, epsilon, pow_2_scales, noop_ptr, dbias_workspace); if (dbias_workspace != nullptr) { const ShapeRepresentation shape_rep = info.same_both_dims ? ShapeRepresentation::SAME_BOTH_DIMS diff --git a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh index 1e32e7464d..b4ff935b8a 100644 --- a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh @@ -135,11 +135,9 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) __shared__ float subamax_colwise_buff[SUBAMAX_BUFF_DIM_Y][CHUNK_DIM_X]; extern __shared__ char dynamic_shmem[]; - uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); // Manually align dynamic SHMEM per TMA requirements using padding // __align__(128) Does not guarantee the pointer to be aligned! - uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) & - ~(static_cast(TMA_SHMEM_ALIGNMENT - 1)); + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); constexpr size_t buff_elems = BUFF_DIM_Y * BUFF_DIM_X; constexpr size_t buff_elems_total = BUFFS_NUM * buff_elems; diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index 615fe47bf7..60c683b440 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -87,7 +87,7 @@ constexpr uint THREADS_PER_BANK = TOTAL_BANKS_WIDTH / SCALE_DIM_X; // 4 = 128 / template + bool WITH_GEMM_SWIZZLED_SCALES, bool kIs2DBlockScaling> __device__ __forceinline__ void process_colwise_stage( const size_t buff, const int stage, const size_t tid_X_colwise, const size_t scales_offset_Y_colwise, const size_t scales_offset_X_colwise, @@ -171,9 +171,13 @@ __device__ __forceinline__ void process_colwise_stage( "+r"(reinterpret_cast(thread_amax_2x)) : "r"(src_smem_ptr), "r"(IN_SHMEM_STRIDE)); } - const float thread_amax = + float thread_amax = static_cast(__hmax(__habs(thread_amax_2x.x), __habs(thread_amax_2x.y))); + if constexpr (kIs2DBlockScaling) { + thread_amax = warp_reduce_max_broadcast(thread_amax); + } + const e8m0_t biased_exponent = ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); // OOB padded region needs to be zeroed out. @@ -246,6 +250,10 @@ __device__ __forceinline__ void process_colwise_stage( } } + if constexpr (kIs2DBlockScaling) { + thread_amax = warp_reduce_max_broadcast(thread_amax); + } + const e8m0_t biased_exponent = ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); // OOB padded region needs to be zeroed out. @@ -270,7 +278,7 @@ __device__ __forceinline__ void process_colwise_stage( template + bool WITH_GEMM_SWIZZLED_SCALES, bool kIs2DBlockScaling> __device__ __forceinline__ void process_rowwise_stage( const size_t buff, const size_t stage_offset_Y, const size_t thread_offset_Y_rowwise, const size_t thread_offset_X_rowwise, const int bank_group, @@ -298,6 +306,8 @@ __device__ __forceinline__ void process_rowwise_stage( auto &sOutRowwise = *reinterpret_cast(sOutRowwise_ptr); const size_t i = thread_offset_Y_rowwise; + const size_t tid_Y_rowwise = thread_offset_Y_rowwise; + const size_t tid_X_rowwise = thread_offset_X_rowwise / SCALE_DIM_X; float thread_amax = 0.0f; float rInCompute[SCALE_DIM_X]; @@ -396,8 +406,31 @@ __device__ __forceinline__ void process_rowwise_stage( } } - const e8m0_t biased_exponent = - ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); + e8m0_t biased_exponent; + if constexpr (kIs2DBlockScaling) { + using AMax2DType = std::conditional_t; + __shared__ e8m0_t block_scales_2d[THREADS_X]; + __shared__ AMax2DType block_amax_2d[THREADS_X * THREADS_Y]; + block_amax_2d[tid_X_rowwise * THREADS_Y + tid_Y_rowwise] = static_cast(thread_amax); + __syncthreads(); + if (tid_Y_rowwise == 0) { + AMax2DType amax_2d = static_cast(0.0f); +#pragma unroll + for (int i = 0; i < THREADS_Y; ++i) { + if constexpr (std::is_same_v) { + amax_2d = fmaxf(amax_2d, block_amax_2d[tid_X_rowwise * THREADS_Y + i]); + } else { + amax_2d = __hmax(amax_2d, block_amax_2d[tid_X_rowwise * THREADS_Y + i]); + } + } + block_scales_2d[tid_X_rowwise] = + ptx::float_to_e8m0(static_cast(amax_2d) * Quantized_Limits::max_norm_rcp); + } + __syncthreads(); + biased_exponent = block_scales_2d[tid_X_rowwise]; + } else { + biased_exponent = ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); + } const size_t stage_scales_offset_Y = scales_offset_Y_rowwise + stage_offset_Y; const size_t stage_scales_offset_X = scales_offset_X_rowwise; @@ -456,7 +489,8 @@ __device__ __forceinline__ void process_rowwise_stage( template + ScalingType SCALING_TYPE, bool WITH_GEMM_SWIZZLED_SCALES, bool kIs2DBlockScaling, + ShapeRepresentation SHAPE_REP> __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel( const __grid_constant__ CUtensorMap tensor_map_input_static, const __grid_constant__ CUtensorMap tensor_map_act_input_static, @@ -472,7 +506,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel constexpr bool COMPUTE_ACTIVATIONS = IS_DACT || IS_ACT; constexpr bool NO_ACTIVATIONS = !COMPUTE_ACTIVATIONS; - if constexpr (NO_ACTIVATIONS) { + if constexpr (NO_ACTIVATIONS && !IS_DBIAS) { if (noop != nullptr && noop[0] == 1.0f) { return; } @@ -580,9 +614,18 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel const size_t scale_stride_colwise = DIVUP_TO_MULTIPLE(cols, scale_alignment_X_colwise); const size_t tensor_base = current_block.tensor_base; - const size_t tensor_base_for_scales = (is_single_tensor && num_tensors > 1) - ? static_cast(offsets_ptr[tensor_id]) - : tensor_base; + size_t tensor_base_for_scales = tensor_base; + size_t tensor_rows_for_scales = rows; + if constexpr (WITH_GEMM_SWIZZLED_SCALES && SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS) { + // The payload is one tall tensor, but GEMM scales are swizzled independently per member. + // Uniform groups omit first_dims/tensor_offsets, so derive the member geometry directly. + tensor_rows_for_scales = first_logical_dim / num_tensors; + tensor_base_for_scales = tensor_id * tensor_rows_for_scales * cols; + } + if constexpr (WITH_GEMM_SWIZZLED_SCALES && + SHAPE_REP == ShapeRepresentation::VARYING_FIRST_DIM) { + tensor_base_for_scales = static_cast(offsets_ptr[tensor_id]); + } const size_t block_id_Y = current_block.block_id_Y; const size_t block_id_X = current_block.block_id_X; const size_t block_offset_Y = current_block.block_offset_Y; @@ -686,15 +729,15 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel const size_t buff = buff_in; if constexpr (COLWISE_SCALING) { process_colwise_stage( + WITH_GEMM_SWIZZLED_SCALES, kIs2DBlockScaling>( buff, stage, tid_X_colwise, scales_offset_Y_colwise, scales_offset_X_colwise, - scale_stride_colwise, tensor_base_for_scales, rows, cols, sIn_ptr, sActIn_ptr, - sCachedAct_ptr, sOutColwise_ptr, scales_colwise, partial_dbias_colwise); + scale_stride_colwise, tensor_base_for_scales, tensor_rows_for_scales, cols, sIn_ptr, + sActIn_ptr, sCachedAct_ptr, sOutColwise_ptr, scales_colwise, partial_dbias_colwise); } if constexpr (ROWWISE_SCALING) { process_rowwise_stage( + WITH_GEMM_SWIZZLED_SCALES, kIs2DBlockScaling>( buff, stage_offset_Y, thread_offset_Y_rowwise, thread_offset_X_rowwise, bank_group, scales_offset_Y_rowwise, scales_offset_X_rowwise, scale_stride_rowwise, rowwise_scale_is_within_bounds, cols, sIn_ptr, sActIn_ptr, sCachedAct_ptr, @@ -954,6 +997,7 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations const size_t block_size = THREADS_PER_CHUNK; const bool with_gemm_swizzled_scales = output->with_gemm_swizzled_scales; + const bool use_2d_quantization = quant_config != nullptr && quant_config->mxfp8_2d_quantization; // Logical shape of a tensor with varying all dims is [1, M*K] if (shape_rep != ShapeRepresentation::VARYING_BOTH_DIMS) { @@ -1090,20 +1134,25 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations last_dims_ptr, use_rowwise_scaling, use_colwise_scaling, IS_DACT); } - auto kernel = - group_quantize_mxfp8_kernel; - - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - - kernel<<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - tensor_map_output_colwise, num_tensors, first_logical_dim, - last_logical_dim, offsets_ptr, first_dims_ptr, last_dims_ptr, - scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, workspace_ptr, - amax_ptr, work_blocks_X, work_blocks_Y); + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_2d_quantization, kIs2DBlockScaling, { + auto kernel = + group_quantize_mxfp8_kernel; + + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + dshmem_size)); + + kernel<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + tensor_map_output_colwise, num_tensors, first_logical_dim, + last_logical_dim, offsets_ptr, first_dims_ptr, last_dims_ptr, + scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, workspace_ptr, + amax_ptr, work_blocks_X, work_blocks_Y); + }); if constexpr (IS_DBIAS) { common::grouped_reduce_dbias( diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index 5c49494e4a..d8d99386bf 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -58,7 +58,7 @@ constexpr size_t THREADS_PER_BANK = TOTAL_BANKS_WIDTH / SCALE_DIM_X; // 4 = 128 template + size_t CHUNK_DIM_X, size_t THREADS_PER_CHUNK, bool kIs2DBlockScaling> __global__ void __launch_bounds__(THREADS_PER_CHUNK) quantize_mxfp8_kernel(const __grid_constant__ CUtensorMap tensor_map_input, const __grid_constant__ CUtensorMap tensor_map_act_input, @@ -77,7 +77,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) using transformer_engine::dispatch::mxfp8::swizzle::gemm_swizzled_scale_idx; - if constexpr (NO_ACTIVATIONS) { + if constexpr (NO_ACTIVATIONS && !IS_DBIAS) { if (noop != nullptr && noop[0] == 1.0f) { return; } @@ -143,11 +143,9 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) constexpr size_t out_mem_rowwise = (ROWWISE_SCALING ? buff_size_aligned_out : 0); extern __shared__ char dynamic_shmem[]; - uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); // Manually align dynamic SHMEM per TMA requirements using padding // __align__(128) Does not guarantee the pointer to be aligned! - uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) & - ~(static_cast(TMA_SHMEM_ALIGNMENT - 1)); + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned IType *in_sh = reinterpret_cast(dshmem); @@ -277,6 +275,10 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) } } + if constexpr (kIs2DBlockScaling) { + thread_amax = warp_reduce_max_broadcast(thread_amax); + } + // 2. Compute E8M0 scaling factor const e8m0_t biased_exponent = ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); @@ -428,8 +430,34 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) } // 2. Compute E8M0 scaling factor - const e8m0_t biased_exponent = - ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); + e8m0_t biased_exponent; + if constexpr (kIs2DBlockScaling) { + using AMax2DType = + std::conditional_t), + IType, float>; + __shared__ e8m0_t block_scales_2d[THREADS_X]; + __shared__ AMax2DType block_amax_2d[THREADS_X * THREADS_Y]; + block_amax_2d[tid_X_rowwise * THREADS_Y + tid_Y_rowwise] = + static_cast(thread_amax); + __syncthreads(); + if (tid_Y_rowwise == 0) { + AMax2DType amax_2d = static_cast(0.0f); +#pragma unroll + for (int i = 0; i < THREADS_Y; ++i) { + if constexpr (std::is_same_v) { + amax_2d = fmaxf(amax_2d, block_amax_2d[tid_X_rowwise * THREADS_Y + i]); + } else { + amax_2d = __hmax(amax_2d, block_amax_2d[tid_X_rowwise * THREADS_Y + i]); + } + } + block_scales_2d[tid_X_rowwise] = ptx::float_to_e8m0( + static_cast(amax_2d) * Quantized_Limits::max_norm_rcp); + } + __syncthreads(); + biased_exponent = block_scales_2d[tid_X_rowwise]; + } else { + biased_exponent = ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); + } const int stage_scales_offset_Y = scales_offset_Y_rowwise + stage_offset_Y; const int stage_scales_offset_X = scales_offset_X_rowwise; size_t scale_idx; @@ -585,7 +613,8 @@ static __global__ void __launch_bounds__(256) template void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, // TODO (ksivamani) - Tensor *output, Tensor *dbias, Tensor *workspace, cudaStream_t stream) { + Tensor *output, Tensor *dbias, Tensor *workspace, const bool use_2d_quantization, + cudaStream_t stream) { using namespace quantize_kernel; #ifndef __HIP_PLATFORM_AMD__ checkCuDriverContext(stream); @@ -740,7 +769,36 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, float *const workspace_ptr = IS_DBIAS ? reinterpret_cast(workspace->data.dptr) : nullptr; float *const amax_ptr = reinterpret_cast(output->amax.dptr); - const float *noop_ptr = reinterpret_cast(noop->data.dptr); + constexpr bool NO_ACTIVATIONS = !(IS_DACT || IS_ACT); + // zero_scales_kernel should ignore the noop tensor whenever the quantization kernel also ignores it, + // which happens when there are no fusions (NO ACT and DBIAS). Since zero_scales_kernel is not templated with these variants + // it doesn't know when to ignore, so we need to override the noop pointer to nullptr before passing it to the kernel. + const float *const noop_ptr = + (NO_ACTIVATIONS && !IS_DBIAS) ? reinterpret_cast(noop->data.dptr) : nullptr; + + // Clear padding before either the generic or specialized kernel writes + // directly into the GEMM-swizzled scale layout. + if (with_gemm_swizzled_scales && (cols % 128 != 0 || rows % 128 != 0)) { + constexpr size_t zero_threads = 256; + if (use_rowwise_scaling) { + const size_t size_bytes = output->scale_inv.buffer_size_bytes(); + if (size_bytes > 0) { + const size_t zero_blocks = DIVUP(size_bytes, zero_threads); + zero_scales_kernel<<>>( + reinterpret_cast(output->scale_inv.dptr), size_bytes, noop_ptr); + NVTE_CHECK_CUDA(cudaGetLastError()); + } + } + if (use_colwise_scaling) { + const size_t size_bytes = output->columnwise_scale_inv.buffer_size_bytes(); + if (size_bytes > 0) { + const size_t zero_blocks = DIVUP(size_bytes, zero_threads); + zero_scales_kernel<<>>( + reinterpret_cast(output->columnwise_scale_inv.dptr), size_bytes, noop_ptr); + NVTE_CHECK_CUDA(cudaGetLastError()); + } + } + } TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( input.dtype(), IType, @@ -763,17 +821,26 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, bidimensional_traits::blockDIM::M) <= max_grid_dim_y; const bool is_full_rowwise_chunk = (cols % 128 == 0); + const bool has_full_bidimensional_chunks = + (rows % bidimensional_traits::colChunkElems == 0) && + (cols % bidimensional_traits::rowChunkElems == 0); + // Both rowwise and bidimensional cast-only kernels select their + // scale layout from WITH_GEMM_SWIZZLED_SCALES. const bool scaling_type_has_specialized_support = (scaling_type == ScalingType::ROWWISE && is_full_rowwise_chunk && rowwise_specialized_grid_fits) || - (scaling_type == ScalingType::BIDIMENSIONAL && + (scaling_type == ScalingType::BIDIMENSIONAL && has_full_bidimensional_chunks && bidimensional_specialized_grid_fits); - if (specialized::hasSpec() && - !WITH_GEMM_SWIZZLED_SCALES && scaling_type_has_specialized_support) { + // Specialized cast-only kernels do not consume the device noop flag. + // Preserve cached outputs by keeping noop-aware calls on the generic path. + if (noop_ptr == nullptr && + specialized::hasSpec() && + !use_2d_quantization && scaling_type_has_specialized_support) { switch (scaling_type) { case ScalingType::ROWWISE: { - using traits = specialized::CastTraits; + using traits = specialized::CastTraits; auto kernel = specialized::quantize_mxfp8_kernel_cast_only; NVTE_CHECK_CUDA(cudaFuncSetAttribute( @@ -786,12 +853,17 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, kernel<<>>( reinterpret_cast(input.data.dptr), reinterpret_cast(output->data.dptr), - scales_rowwise_ptr, rows, cols, scale_stride_rowwise, scale_stride_colwise); + scales_rowwise_ptr, noop_ptr, rows, cols, scale_stride_rowwise, + scale_stride_colwise); break; } case ScalingType::BIDIMENSIONAL: { - using traits = specialized::CastTraits; + using traits = + specialized::CastTraitsSwizzle; auto kernel = specialized::quantize_mxfp8_kernel_cast_only; NVTE_CHECK_CUDA(cudaFuncSetAttribute( @@ -823,8 +895,8 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, (rows + traits::blockDIM::M - 1) / traits::blockDIM::M); kernel<<>>( tensor_map_input, tensor_map_rowwise_output, tensor_map_colwise_output, - scales_rowwise_ptr, scales_colwise_ptr, rows, cols, scale_stride_rowwise, - scale_stride_colwise); + scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, rows, cols, + scale_stride_rowwise, scale_stride_colwise); break; } @@ -881,79 +953,56 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, const size_t dshmem_size = in_mem + out_mem + TMA_SHMEM_ALIGNMENT; - // Zero out swizzled scales if padding is needed - /// TODO (tmoon) Handle this within the cast kernel - if (with_gemm_swizzled_scales) { - constexpr size_t TILE_DIM_X = 128; // Tile dim in data buffer - constexpr size_t TILE_DIM_Y = 128; - if (cols % TILE_DIM_X != 0 || rows % TILE_DIM_Y != 0) { - // Use a noop-aware zero kernel so that the clear is skipped - // when quantization is a noop (e.g. FP8 weight caching). - constexpr size_t zero_threads = 256; - if (use_rowwise_scaling) { - const size_t size_bytes = output->scale_inv.buffer_size_bytes(); - if (size_bytes > 0) { - const size_t zero_blocks = DIVUP(size_bytes, zero_threads); - zero_scales_kernel<<>>( - reinterpret_cast(output->scale_inv.dptr), size_bytes, - noop_ptr); - NVTE_CHECK_CUDA(cudaGetLastError()); - } - } - if (use_colwise_scaling) { - const size_t size_bytes = output->columnwise_scale_inv.buffer_size_bytes(); - if (size_bytes > 0) { - const size_t zero_blocks = DIVUP(size_bytes, zero_threads); - zero_scales_kernel<<>>( - reinterpret_cast(output->columnwise_scale_inv.dptr), - size_bytes, noop_ptr); - NVTE_CHECK_CUDA(cudaGetLastError()); - } - } - } - } - switch (scaling_type) { case ScalingType::ROWWISE: { - auto kernel = quantize_mxfp8_kernel; - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - - kernel<<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, - workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, - scale_stride_colwise); + TRANSFORMER_ENGINE_SWITCH_CONDITION(use_2d_quantization, kIs2DBlockScaling, { + auto kernel = + quantize_mxfp8_kernel; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + + kernel<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, + workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, + scale_stride_colwise); + }); break; } case ScalingType::COLWISE: { - auto kernel = quantize_mxfp8_kernel; - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - - kernel<<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, - workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, - scale_stride_colwise); + TRANSFORMER_ENGINE_SWITCH_CONDITION(use_2d_quantization, kIs2DBlockScaling, { + auto kernel = + quantize_mxfp8_kernel; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + + kernel<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, + workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, + scale_stride_colwise); + }); break; } case ScalingType::BIDIMENSIONAL: { - auto kernel = quantize_mxfp8_kernel; - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - - kernel<<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, - workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, - scale_stride_colwise); + TRANSFORMER_ENGINE_SWITCH_CONDITION(use_2d_quantization, kIs2DBlockScaling, { + auto kernel = + quantize_mxfp8_kernel; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + + kernel<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, + workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, + scale_stride_colwise); + }); break; } } NVTE_CHECK_CUDA(cudaGetLastError()); diff --git a/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh index 9459f0273a..507122cf8d 100644 --- a/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh @@ -14,8 +14,9 @@ #include #include "../../../util/ptx.cuh" +#include "../swizzle.cuh" // gemm_swizzled_scale_idx (parent dir, GEMM scale swizzle) #include "state_counter.cuh" -#include "swizzle.cuh" +#include "swizzle.cuh" // specialized/swizzle.cuh (TMA input bank-conflict swizzle) namespace transformer_engine { namespace dispatch { @@ -24,6 +25,10 @@ namespace quantize_kernel { namespace specialized { namespace ptx = transformer_engine::ptx; + +// Bring in the GEMM-swizzled scale index helper (from ../swizzle.cuh). +// Used only when the kernel is instantiated with CastTraits::_with_swizzled_scales=true. +using transformer_engine::dispatch::mxfp8::swizzle::gemm_swizzled_scale_idx; namespace { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) @@ -122,19 +127,20 @@ struct Layout { static constexpr int32_t num = M * N; }; -template +template struct CastTraits; // 1x32 -template -struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/false> { +template +struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/false, _kSwizzled> { static constexpr bool isRowwise = true; static constexpr bool isColwise = false; using IType = _IType; using OType = _OType; static constexpr int32_t chunkElems = 32; - using threadLayout = Layout<1, 32>; + using threadLayout = Layout<1, THREADS_PER_WARP>; static constexpr int32_t numThreadsPerChunk = 1; static constexpr int32_t warpDimM = threadLayout::M; static constexpr int32_t warpDimN = threadLayout::N * chunkElems; @@ -151,14 +157,22 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/false> { using iterLayout = Layout<1, 1>; static constexpr int32_t blockDimM = iterLayout::M * blockIterDimM; static constexpr int32_t blockDimN = iterLayout::N * blockIterDimN; + static constexpr int32_t rowwiseScaleStride = blockDimN / chunkElems; + using PreferredDataType = std::conditional_t< + rowwiseScaleStride % 16 == 0, uint4, + std::conditional_t< + rowwiseScaleStride % 8 == 0, uint2, + std::conditional_t>>>; static constexpr int32_t numStages = 1; static constexpr int32_t numPrefetch = numStages - 1; static constexpr bool _use_cvt_4x = true; static constexpr bool _cache_rowwise_scale_in_smem = true; + static constexpr bool _with_swizzled_scales = _kSwizzled; - static constexpr int32_t numThreads = warpLayout::num * 32; + static constexpr int32_t numThreads = warpLayout::num * THREADS_PER_WARP; static constexpr size_t smem_rowwise_scale = _cache_rowwise_scale_in_smem ? (blockDimM * (blockDimN / chunkElems) * sizeof(e8m0_t)) : 0ul; @@ -170,10 +184,14 @@ template = 0> __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__restrict__ input, typename CastTraits::OType *__restrict__ output, - e8m0_t *__restrict__ scales_rowwise, int32_t rows, - int32_t cols, int32_t scale_stride_rowwise, + e8m0_t *__restrict__ scales_rowwise, + const float *noop, int32_t rows, int32_t cols, + int32_t scale_stride_rowwise, int32_t scale_stride_colwise) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + if (noop != nullptr && noop[0] == 1.0f) { + return; + } using IType = typename CastTraits::IType; using OType = typename CastTraits::OType; using inputUnitType = typename CastTraits::inputUnitType; @@ -494,13 +512,8 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re block_coords.y = blockIdx.y * CastTraits::blockDimM; block_coords.x = blockIdx.x * CastTraits::blockDimN; - constexpr int32_t stride_in_smem = CastTraits::blockDimN / CastTraits::chunkElems; - using PreferredDataType = std::conditional_t< - stride_in_smem % 16 == 0, uint4, - std::conditional_t< - stride_in_smem % 8 == 0, uint2, - std::conditional_t>>>; + constexpr int32_t stride_in_smem = CastTraits::rowwiseScaleStride; + using PreferredDataType = typename CastTraits::PreferredDataType; int2 end_coords; end_coords.y = std::min(block_coords.y + CastTraits::blockDimM, rows); @@ -510,7 +523,40 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re valid_coords.y = end_coords.y - block_coords.y; valid_coords.x = end_coords.x - (block_coords.x / CastTraits::chunkElems); - if (scale_stride_rowwise % sizeof(PreferredDataType) != 0) { + if constexpr (CastTraits::_with_swizzled_scales) { + // Four adjacent rowwise scale columns are contiguous in the GEMM scale + // layout, so write them as one uint32_t whenever possible. + constexpr int32_t cols_per_group = 4; + const int32_t groups_per_row = valid_coords.x / cols_per_group; + const int32_t total_groups = valid_coords.y * groups_per_row; + const int32_t base_col = block_coords.x / CastTraits::chunkElems; + const size_t num_tiles_x = DIVUP(cols, static_cast(128)); + + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_groups; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + const int32_t row = i / groups_per_row; + const int32_t col = (i % groups_per_row) * cols_per_group; + const uint32_t value = + *reinterpret_cast(&sRowwiseScale[row * stride_in_smem + col]); + const size_t idx = + gemm_swizzled_scale_idx(block_coords.y + row, base_col + col, num_tiles_x); + *reinterpret_cast(&scales_rowwise[idx]) = value; + } + + const int32_t remaining_start = groups_per_row * cols_per_group; + const int32_t remaining_per_row = valid_coords.x - remaining_start; + if (remaining_per_row > 0) { + const int32_t total_remaining = valid_coords.y * remaining_per_row; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_remaining; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + const int32_t row = i / remaining_per_row; + const int32_t col = remaining_start + (i % remaining_per_row); + const size_t idx = + gemm_swizzled_scale_idx(block_coords.y + row, base_col + col, num_tiles_x); + scales_rowwise[idx] = sRowwiseScale[row * stride_in_smem + col]; + } + } + } else if (scale_stride_rowwise % sizeof(PreferredDataType) != 0) { using DataType = int32_t; constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; @@ -523,8 +569,9 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + block_coords.x / CastTraits::chunkElems); - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { int32_t row = i / num_threads_per_row; int32_t col = i % num_threads_per_row; gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; @@ -542,8 +589,9 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + block_coords.x / CastTraits::chunkElems); - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { int32_t row = i / num_threads_per_row; int32_t col = i % num_threads_per_row; gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; @@ -573,11 +621,12 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { static constexpr int32_t rowChunkElems = 32; static constexpr int32_t colChunkElems = 32; - using rowThreadLayout = Layout<32, 1>; // 32x1 + using rowThreadLayout = Layout; // 32x1 using colThreadLayout = Layout; // 1x32 static_assert(rowThreadLayout::num == colThreadLayout::num, "rowThreadLayout::num must be equal to colThreadLayout::num"); - static_assert(rowThreadLayout::num == 32, "rowThreadLayout::num must be 32"); + static_assert(rowThreadLayout::num == THREADS_PER_WARP, + "rowThreadLayout::num must match the warp size"); using rowWarpDim = Layout; using colWarpDim = Layout; @@ -599,6 +648,13 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { using iterLayout = Layout<1, 4>; using blockDIM = Layout; + static constexpr int32_t rowwiseScaleStride = blockDIM::N / rowChunkElems; + using PreferredDataType = std::conditional_t< + rowwiseScaleStride % 16 == 0, uint4, + std::conditional_t< + rowwiseScaleStride % 8 == 0, uint2, + std::conditional_t>>>; static constexpr int32_t numStages = 2; @@ -632,7 +688,7 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { "It requires aligned smem pointer"); static constexpr int32_t numWarps = warpLayout::num + 2 * (int32_t)_use_warp_specialization; - static constexpr int32_t numThreads = numWarps * 32; + static constexpr int32_t numThreads = numWarps * THREADS_PER_WARP; static_assert(numThreads <= 1024, "numThreads must be less than or equal to 1024"); static constexpr size_t smemInputPerWarp = warpDim::num * sizeof(IType); @@ -656,7 +712,9 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { static constexpr bool _need_smem_for_colwise_reduce = _colwise_source_coming_from_rowwise; // && _colwise_reduce_max != ColwiseReduceMax::Redux; static constexpr size_t smem_colwise_reduce = - _need_smem_for_colwise_reduce ? 32 * warpLayout::num * sizeof(ColwiseReduceDataType) : 0ul; + _need_smem_for_colwise_reduce + ? THREADS_PER_WARP * warpLayout::num * sizeof(ColwiseReduceDataType) + : 0ul; static constexpr size_t smem_alignment = _tma_swizzle ? 1024ul : 128ul; static constexpr size_t smem = _reuse_input_out_smem @@ -666,9 +724,138 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { smem_alignment + smem_rowwise_scale + smem_colwise_reduce); }; -__device__ __forceinline__ intptr_t align_to(intptr_t x, intptr_t align) { - return (x + align - 1) & ~((align)-1); -} +// Standalone trait for the non-warp-specialized rowwise+colwise cast_only kernel. +// Exposes numStages, iterM, iterN, and the two colwise-scale features as +// caller-controllable template axes. Both colwise flags default to true, +// giving callers the swizzled + colwise-scale-cached fast path by default. +// +// This trait duck-types the same interface CastTraits<_, _, true, true> exposes, +// so it drops into quantize_mxfp8_kernel_cast_only without any +// changes to the kernel signature. Kernel #1 (rowwise-only) and Kernel #2 +// (warp-specialized row+col) don't accept it: kernel #1 requires isColwise=false, +// kernel #2 requires _use_warp_specialization=true - both are wrong here. +template +struct CastTraitsSwizzle { + static constexpr bool isRowwise = true; + static constexpr bool isColwise = true; + using IType = _IType; + using OType = _OType; + + static constexpr int32_t rowChunkElems = 32; + static constexpr int32_t colChunkElems = 32; + + using rowThreadLayout = Layout; // 32x1 + using colThreadLayout = Layout; // 1x32 + static_assert(rowThreadLayout::num == colThreadLayout::num, + "rowThreadLayout::num must be equal to colThreadLayout::num"); + static_assert(rowThreadLayout::num == THREADS_PER_WARP, + "rowThreadLayout::num must match the warp size"); + + using rowWarpDim = Layout; + using colWarpDim = Layout; + using warpDim = + Layout; + + static constexpr bool _tma_swizzle = true; + using warpLayout = Layout<1, 2>; + static_assert(_tma_swizzle ? (warpLayout::N == 2) : true); + static constexpr CUtensorMapSwizzle input_swizzle_pattern = + _tma_swizzle ? CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B + : CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE; + + static constexpr CUtensorMapSwizzle output_swizzle_pattern = + _tma_swizzle ? CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_64B + : CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE; + + using blockIterDim = Layout; + + using iterLayout = Layout<_IterM, _IterN>; + using blockDIM = Layout; + static constexpr int32_t rowwiseScaleStride = blockDIM::N / rowChunkElems; + using PreferredDataType = std::conditional_t< + rowwiseScaleStride % 16 == 0, uint4, + std::conditional_t< + rowwiseScaleStride % 8 == 0, uint2, + std::conditional_t>>>; + + static constexpr int32_t numStages = _NumStages; + + using inputUnitType = uint4; + static constexpr int32_t rowNumElemsPerUnit = sizeof(inputUnitType) / sizeof(IType); + static constexpr int32_t rowNumUnitsPerChunk = rowChunkElems / rowNumElemsPerUnit; + using inputElemSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<3, 3, 3>, swz::Linear>; + using inputUnitSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<3, 0, 3>, swz::Linear>; + + using colIndexSwz = swz::Swizzle<5, 0, 5>; + + using rowOutputUnitType = uint4; + static constexpr int32_t rowNumOutUnitsPerChunk = + rowChunkElems * sizeof(OType) / sizeof(rowOutputUnitType); + static constexpr int32_t rowOutNumElemsPerUnit = sizeof(rowOutputUnitType) / sizeof(OType); + + using rowOutputChunkSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<2, 0, 3>, swz::Linear>; + using colOutputSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<2, 4, 3>, swz::Linear>; + + static constexpr bool _use_cvt_4x = true; + static constexpr bool _use_warp_specialization = false; + static constexpr bool _need_wait_group = iterLayout::num > numStages; + static constexpr bool _reuse_input_out_smem = false; + static_assert(_reuse_input_out_smem == false, "Just don't use it"); + static constexpr bool _cache_rowwise_scale_in_smem = true; + + static constexpr bool _colwise_source_coming_from_rowwise = true; + static constexpr ColwiseReduceMax _colwise_reduce_max = ColwiseReduceMax::Redux; + static_assert(_colwise_reduce_max != ColwiseReduceMax::RedAsync, + "It requires aligned smem pointer"); + + // The two colwise-scale features exposed as caller-controllable trait axes. + // Both default to true so callers get the vectorized swizzled path by default. + static constexpr bool _cache_colwise_scale_in_smem = _kCacheColwise; + static constexpr bool _with_swizzled_scales = _kSwizzled; + + static constexpr int32_t numWarps = warpLayout::num + 2 * (int32_t)_use_warp_specialization; + static constexpr int32_t numThreads = numWarps * THREADS_PER_WARP; + static_assert(numThreads <= 1024, "numThreads must be less than or equal to 1024"); + + static constexpr size_t smemInputPerWarp = warpDim::num * sizeof(IType); + static constexpr size_t smemInputPerBlock = smemInputPerWarp * warpLayout::num; + + static constexpr size_t smemRowwiseOutputPerWarp = warpDim::num * sizeof(OType); + static constexpr size_t smemRowwiseOutputPerBlock = smemRowwiseOutputPerWarp * warpLayout::num; + + static constexpr size_t smemColwiseOutputPerWarp = warpDim::num * sizeof(OType); + static constexpr size_t smemColwiseOutputPerBlock = smemColwiseOutputPerWarp * warpLayout::num; + + static constexpr size_t smemInput = smemInputPerBlock * numStages; + static constexpr size_t smemRowwiseOutput = smemRowwiseOutputPerBlock * numStages; + static constexpr size_t smemColwiseOutput = smemColwiseOutputPerBlock * numStages; + + static constexpr size_t smem_rowwise_scale = + _cache_rowwise_scale_in_smem ? (blockDIM::M * (blockDIM::N / rowChunkElems) * sizeof(e8m0_t)) + : 0ul; + + // Extra shmem for cached colwise scales - only when the flag is on. + static constexpr size_t smem_colwise_scale = + _cache_colwise_scale_in_smem ? (blockDIM::M / colChunkElems) * blockDIM::N * sizeof(e8m0_t) + : 0ul; + + using ColwiseReduceDataType = float; + static constexpr bool _need_smem_for_colwise_reduce = _colwise_source_coming_from_rowwise; + static constexpr size_t smem_colwise_reduce = + _need_smem_for_colwise_reduce + ? THREADS_PER_WARP * warpLayout::num * sizeof(ColwiseReduceDataType) + : 0ul; + + static constexpr size_t smem_alignment = _tma_swizzle ? 1024ul : 128ul; + static constexpr size_t smem = + _reuse_input_out_smem + ? (std::max(smemInput, smemColwiseOutput) + smemRowwiseOutput + smem_alignment + + smem_rowwise_scale + smem_colwise_scale + smem_colwise_reduce) + : (smemInput + smemRowwiseOutput + smemColwiseOutput + smem_alignment + + smem_rowwise_scale + smem_colwise_scale + smem_colwise_reduce); +}; // 32x32 template = 1000) + if (noop != nullptr && noop[0] == 1.0f) { + return; + } using IType = typename CastTraits::IType; using OType = typename CastTraits::OType; using inputUnitType = typename CastTraits::inputUnitType; @@ -699,8 +889,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( block_coords.x = blockIdx.x * CastTraits::blockDIM::N; extern __shared__ char smem[]; - char *smemAligned = reinterpret_cast( - align_to(reinterpret_cast(smem), CastTraits::smem_alignment)); + char *smemAligned = align_up(smem, CastTraits::smem_alignment); IType *sInput = reinterpret_cast(smemAligned); inputUnitType *sInputUnit = reinterpret_cast(sInput); @@ -721,12 +910,12 @@ __global__ void quantize_mxfp8_kernel_cast_only( if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t)); - sColwiseReduce += warpId * 32; + sColwiseReduce += warpId * THREADS_PER_WARP; } } else if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( sColOutput + CastTraits::blockIterDim::num * CastTraits::numStages); - sColwiseReduce += warpId * 32; + sColwiseReduce += warpId * THREADS_PER_WARP; } // TODO: maybe we can assign a different barrier for each warp @@ -737,8 +926,8 @@ __global__ void quantize_mxfp8_kernel_cast_only( #pragma unroll for (int32_t i = 0; i < CastTraits::numStages; i++) { ptx::mbarrier_init(&ldg_producer[i], 1); - ptx::mbarrier_init(&ldg_consumer[i], CastTraits::warpLayout::num * 32); - ptx::mbarrier_init(&stg_producer[i], CastTraits::warpLayout::num * 32); + ptx::mbarrier_init(&ldg_consumer[i], CastTraits::warpLayout::num * THREADS_PER_WARP); + ptx::mbarrier_init(&stg_producer[i], CastTraits::warpLayout::num * THREADS_PER_WARP); ptx::mbarrier_init(&stg_consumer[i], 1); } ptx::fence_mbarrier_init_release_cluster(); @@ -1078,15 +1267,10 @@ __global__ void quantize_mxfp8_kernel_cast_only( } if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { - ptx::numbered_barrier_sync(CastTraits::warpLayout::num * 32, 0u); + ptx::numbered_barrier_sync(CastTraits::warpLayout::num * THREADS_PER_WARP, 0u); - constexpr int32_t stride_in_smem = CastTraits::blockDIM::N / CastTraits::rowChunkElems; - using PreferredDataType = std::conditional_t< - stride_in_smem % 16 == 0, uint4, - std::conditional_t< - stride_in_smem % 8 == 0, uint2, - std::conditional_t>>>; + constexpr int32_t stride_in_smem = CastTraits::rowwiseScaleStride; + using PreferredDataType = typename CastTraits::PreferredDataType; int2 end_coords; end_coords.y = std::min(block_coords.y + CastTraits::blockDIM::M, rows); @@ -1110,8 +1294,9 @@ __global__ void quantize_mxfp8_kernel_cast_only( reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + block_coords.x / CastTraits::rowChunkElems); - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { int32_t row = i / num_threads_per_row; int32_t col = i % num_threads_per_row; gScales[row * gmem_stride_in_group + col] = @@ -1130,8 +1315,9 @@ __global__ void quantize_mxfp8_kernel_cast_only( reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + block_coords.x / CastTraits::rowChunkElems); - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { int32_t row = i / num_threads_per_row; int32_t col = i % num_threads_per_row; gScales[row * gmem_stride_in_group + col] = @@ -1150,9 +1336,12 @@ __global__ void quantize_mxfp8_kernel_cast_only( const __grid_constant__ CUtensorMap tensor_map_input, const __grid_constant__ CUtensorMap tensor_map_rowwise_output, const __grid_constant__ CUtensorMap tensor_map_colwise_output, e8m0_t *scales_rowwise, - e8m0_t *scales_colwise, int32_t rows, int32_t cols, int32_t scale_stride_rowwise, - int32_t scale_stride_colwise) { + e8m0_t *scales_colwise, const float *noop, int32_t rows, int32_t cols, + int32_t scale_stride_rowwise, int32_t scale_stride_colwise) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + if (noop != nullptr && noop[0] == 1.0f) { + return; + } using IType = typename CastTraits::IType; using OType = typename CastTraits::OType; using inputUnitType = typename CastTraits::inputUnitType; @@ -1169,8 +1358,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( block_coords.x = blockIdx.x * CastTraits::blockDIM::N; extern __shared__ char smem[]; - char *smemAligned = reinterpret_cast( - align_to(reinterpret_cast(smem), CastTraits::smem_alignment)); + char *smemAligned = align_up(smem, CastTraits::smem_alignment); IType *sInput = reinterpret_cast(smemAligned); inputUnitType *sInputUnit = reinterpret_cast(sInput); @@ -1181,15 +1369,20 @@ __global__ void quantize_mxfp8_kernel_cast_only( // colwise output will reuse input buffer OType *sColOutput; e8m0_t *sRowwiseScale = nullptr; + e8m0_t *sColwiseScale = nullptr; ColwiseReduceDataType *sColwiseReduce = nullptr; if constexpr (CastTraits::_reuse_input_out_smem) { sColOutput = reinterpret_cast(sInput); if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { sRowwiseScale = reinterpret_cast(sRowOutput + CastTraits::blockIterDim::num * CastTraits::numStages); + if constexpr (CastTraits::_cache_colwise_scale_in_smem) { + sColwiseScale = sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t); + } if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( - sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t)); + sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t) + + CastTraits::smem_colwise_scale); } } else if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( @@ -1201,9 +1394,13 @@ __global__ void quantize_mxfp8_kernel_cast_only( if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { sRowwiseScale = reinterpret_cast(sColOutput + CastTraits::blockIterDim::num * CastTraits::numStages); + if constexpr (CastTraits::_cache_colwise_scale_in_smem) { + sColwiseScale = sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t); + } if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( - sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t)); + sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t) + + CastTraits::smem_colwise_scale); } } else if constexpr (CastTraits::_need_smem_for_colwise_reduce) { sColwiseReduce = reinterpret_cast( @@ -1213,7 +1410,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( rowOutputUnitType *sColOutputUnit = reinterpret_cast(sColOutput); if constexpr (CastTraits::_need_smem_for_colwise_reduce) { - sColwiseReduce += warpId * 32; + sColwiseReduce += warpId * THREADS_PER_WARP; } __shared__ uint64_t producer[CastTraits::numStages]; @@ -1233,7 +1430,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( } if constexpr (CastTraits::_colwise_source_coming_from_rowwise && CastTraits::_colwise_reduce_max == ColwiseReduceMax::RedAsync) { - ptx::mbarrier_init(colwise_reduce_barrier, 32); + ptx::mbarrier_init(colwise_reduce_barrier, THREADS_PER_WARP); } ptx::fence_mbarrier_init_release_cluster(); @@ -1253,18 +1450,35 @@ __global__ void quantize_mxfp8_kernel_cast_only( (threadIdx.x % CastTraits::rowThreadLayout::N) * (CastTraits::rowChunkElems / CastTraits::rowNumElemsPerUnit); - size_t rowwise_scale_base_offset = - (block_coords.y + warp_coords.y + (threadIdx.x / CastTraits::rowThreadLayout::N)) * - static_cast(scale_stride_rowwise) + + // Scale coordinates in absolute (compact) scale-tensor space. Shared by both + // compact-layout offsets and by the CastTraits::_with_swizzled_scales branches below. + const int32_t row_scale_row_base = + block_coords.y + warp_coords.y + (threadIdx.x / CastTraits::rowThreadLayout::N); + const int32_t row_scale_col_base = (block_coords.x + warp_coords.x + (threadIdx.x % CastTraits::rowThreadLayout::N) * CastTraits::rowChunkElems) / - CastTraits::rowChunkElems; + CastTraits::rowChunkElems; + const int32_t col_scale_row_base = + (block_coords.y + warp_coords.y + + (threadIdx.x / CastTraits::colThreadLayout::N) * CastTraits::colChunkElems) / + CastTraits::colChunkElems; + const int32_t col_scale_col_base = + block_coords.x + warp_coords.x + (threadIdx.x % CastTraits::colThreadLayout::N); + + size_t rowwise_scale_base_offset = + static_cast(row_scale_row_base) * static_cast(scale_stride_rowwise) + + row_scale_col_base; size_t colwise_scale_base_offset = - ((block_coords.y + warp_coords.y + - (threadIdx.x / CastTraits::colThreadLayout::N) * CastTraits::colChunkElems) / - CastTraits::colChunkElems) * - static_cast(scale_stride_colwise) + - (block_coords.x + warp_coords.x + (threadIdx.x % CastTraits::colThreadLayout::N)); + static_cast(col_scale_row_base) * static_cast(scale_stride_colwise) + + col_scale_col_base; + + // Precomputed swizzle constants (each swizzle tile is 128 rows x 4 cols in scale space). + // Rowwise scale tensor has DIVUP(cols, 128) tiles across; + // colwise scale tensor has DIVUP(rows, 128) tiles across (X/Y axes are transposed). + const size_t row_swz_num_tiles_X = + CastTraits::_with_swizzled_scales ? DIVUP(cols, static_cast(128)) : 0; + const size_t col_swz_num_tiles_X = + CastTraits::_with_swizzled_scales ? DIVUP(rows, static_cast(128)) : 0; constexpr int32_t rowwise_scale_stride_in_smem = CastTraits::blockDIM::N / CastTraits::rowChunkElems; @@ -1391,6 +1605,12 @@ __global__ void quantize_mxfp8_kernel_cast_only( iter_m * CastTraits::blockIterDim::M * rowwise_scale_stride_in_smem + iter_n * (CastTraits::blockIterDim::N / CastTraits::rowChunkElems); sRowwiseScale[rowwise_scale_offset] = row_biased_exponent; + } else if constexpr (CastTraits::_with_swizzled_scales) { + int32_t abs_row = row_scale_row_base + iter_m * CastTraits::blockIterDim::M; + int32_t abs_col = row_scale_col_base + + iter_n * (CastTraits::blockIterDim::N / CastTraits::rowChunkElems); + size_t idx = gemm_swizzled_scale_idx(abs_row, abs_col, row_swz_num_tiles_X); + scales_rowwise[idx] = row_biased_exponent; } else { size_t rowwise_scale_offset = rowwise_scale_base_offset + @@ -1406,12 +1626,32 @@ __global__ void quantize_mxfp8_kernel_cast_only( e8m0_t col_biased_exponent = to_e8m0(col_amax); float col_scale_inverse = ptx::exp2f_rcp(col_biased_exponent); sColwiseReduce[threadIdx.x] = col_scale_inverse; - size_t colwise_scale_offset = - colwise_scale_base_offset + - iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems) * - static_cast(scale_stride_colwise) + - iter_n * CastTraits::blockIterDim::N; - scales_colwise[colwise_scale_offset] = col_biased_exponent; + if constexpr (CastTraits::_cache_colwise_scale_in_smem) { + // Cache in shmem; end-of-kernel flush handles gmem indexing. + int32_t smem_row = (warp_coords.y + (threadIdx.x / CastTraits::colThreadLayout::N) * + CastTraits::colChunkElems) / + CastTraits::colChunkElems + + iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems); + int32_t smem_col = warp_coords.x + (threadIdx.x % CastTraits::colThreadLayout::N) + + iter_n * CastTraits::blockIterDim::N; + sColwiseScale[smem_row * CastTraits::blockDIM::N + smem_col] = col_biased_exponent; + } else if constexpr (CastTraits::_with_swizzled_scales) { + int32_t abs_row = col_scale_row_base + + iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems); + int32_t abs_col = col_scale_col_base + iter_n * CastTraits::blockIterDim::N; + // Colwise scale tensor's X/Y axes are transposed vs rowwise + // (see col_swz_num_tiles_X = DIVUP(rows, 128)), so pass + // (abs_col, abs_row) - abs_col is the swizzle "row" dim. + size_t idx = gemm_swizzled_scale_idx(abs_col, abs_row, col_swz_num_tiles_X); + scales_colwise[idx] = col_biased_exponent; + } else { + size_t colwise_scale_offset = + colwise_scale_base_offset + + iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems) * + static_cast(scale_stride_colwise) + + iter_n * CastTraits::blockIterDim::N; + scales_colwise[colwise_scale_offset] = col_biased_exponent; + } __syncwarp(); } } @@ -1536,58 +1776,210 @@ __global__ void quantize_mxfp8_kernel_cast_only( if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { constexpr int32_t stride_in_smem = CastTraits::blockDIM::N / CastTraits::rowChunkElems; - using PreferredDataType = std::conditional_t< - stride_in_smem % 16 == 0, uint4, - std::conditional_t< - stride_in_smem % 8 == 0, uint2, - std::conditional_t>>>; int2 end_coords; end_coords.y = std::min(block_coords.y + CastTraits::blockDIM::M, rows); - end_coords.x = std::min((block_coords.x + CastTraits::blockDIM::N) / CastTraits::rowChunkElems, - scale_stride_rowwise); + if constexpr (CastTraits::_with_swizzled_scales) { + end_coords.x = + std::min((block_coords.x + CastTraits::blockDIM::N) / CastTraits::rowChunkElems, + DIVUP(cols, static_cast(CastTraits::rowChunkElems))); + } else { + // The compact layout's padded entries are consumed by a later swizzle. + end_coords.x = + std::min((block_coords.x + CastTraits::blockDIM::N) / CastTraits::rowChunkElems, + scale_stride_rowwise); + } int2 valid_coords; valid_coords.y = end_coords.y - block_coords.y; valid_coords.x = end_coords.x - (block_coords.x / CastTraits::rowChunkElems); - if (scale_stride_rowwise % sizeof(PreferredDataType) != 0) { - using DataType = int32_t; - constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); - constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; - - int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group); - int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group; - - DataType *sScales = reinterpret_cast(sRowwiseScale); - DataType *gScales = - reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + - block_coords.x / CastTraits::rowChunkElems); + if constexpr (CastTraits::_with_swizzled_scales) { + // Swizzled flush: within a 128x4 swizzle tile, 4 consecutive column entries + // for the same row live at 4 consecutive gmem bytes. Group by 4 so each + // thread writes a uint32_t when col%4 == 0, then scalar tail for remainder. + constexpr int32_t cols_per_group = 4; + const int32_t groups_per_row = valid_coords.x / cols_per_group; + const int32_t total_groups = valid_coords.y * groups_per_row; + const int32_t base_col = block_coords.x / CastTraits::rowChunkElems; + + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_groups; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / groups_per_row; + int32_t group = i % groups_per_row; + int32_t col = group * cols_per_group; + + uint32_t val4 = + *reinterpret_cast(&sRowwiseScale[row * stride_in_smem + col]); + + int32_t abs_row = block_coords.y + row; + int32_t abs_col = base_col + col; + size_t idx = gemm_swizzled_scale_idx(abs_row, abs_col, row_swz_num_tiles_X); + *reinterpret_cast(&scales_rowwise[idx]) = val4; + } - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { - int32_t row = i / num_threads_per_row; - int32_t col = i % num_threads_per_row; - gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; + // Tail (valid_coords.x % 4 != 0) + const int32_t remaining_start = groups_per_row * cols_per_group; + const int32_t remaining_per_row = valid_coords.x - remaining_start; + if (remaining_per_row > 0) { + const int32_t total_remaining = valid_coords.y * remaining_per_row; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_remaining; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / remaining_per_row; + int32_t col = remaining_start + (i % remaining_per_row); + e8m0_t val = sRowwiseScale[row * stride_in_smem + col]; + int32_t abs_row = block_coords.y + row; + int32_t abs_col = base_col + col; + size_t idx = gemm_swizzled_scale_idx(abs_row, abs_col, row_swz_num_tiles_X); + scales_rowwise[idx] = val; + } } } else { - using DataType = PreferredDataType; - constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); - constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; + using PreferredDataType = typename CastTraits::PreferredDataType; - int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group); - int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group; + if (scale_stride_rowwise % sizeof(PreferredDataType) != 0) { + using DataType = int32_t; + constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); + constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; - DataType *sScales = reinterpret_cast(sRowwiseScale); - DataType *gScales = - reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + - block_coords.x / CastTraits::rowChunkElems); + int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group); + int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group; - for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); - i += CastTraits::warpLayout::num * 32) { - int32_t row = i / num_threads_per_row; - int32_t col = i % num_threads_per_row; - gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; + DataType *sScales = reinterpret_cast(sRowwiseScale); + DataType *gScales = + reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + + block_coords.x / CastTraits::rowChunkElems); + + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / num_threads_per_row; + int32_t col = i % num_threads_per_row; + gScales[row * gmem_stride_in_group + col] = + sScales[row * num_groups_per_row_in_smem + col]; + } + } else { + using DataType = PreferredDataType; + constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); + constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; + + int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group); + int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group; + + DataType *sScales = reinterpret_cast(sRowwiseScale); + DataType *gScales = + reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + + block_coords.x / CastTraits::rowChunkElems); + + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; + i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / num_threads_per_row; + int32_t col = i % num_threads_per_row; + gScales[row * gmem_stride_in_group + col] = + sScales[row * num_groups_per_row_in_smem + col]; + } + } + } + } + + // Cached colwise scale flush (swizzled path only). Same barrier semantics as + // the rowwise flush above: the last-iter __syncthreads already ordered every + // in-loop sColwiseScale byte store before this block reads them. + if constexpr (CastTraits::_cache_colwise_scale_in_smem && CastTraits::_with_swizzled_scales) { + const int32_t scale_row_base = block_coords.y / CastTraits::colChunkElems; + // DIVUP so a partial last block (e.g. rows=993, colChunkElems=32 -> last CTA + // has 1 valid input row) still emits its scale row. Truncating divison here + // drops the last partial scale row, causing the fast path to diverge from + // nvte_swizzle_scaling_factors on non-multiple-of-32-rows inputs. + const int32_t valid_rows = + DIVUP(std::min(block_coords.y + CastTraits::blockDIM::M, rows) - block_coords.y, + static_cast(CastTraits::colChunkElems)); + const int32_t valid_cols = + std::min(block_coords.x + CastTraits::blockDIM::N, cols) - block_coords.x; + + // In GEMM swizzle, contiguous bytes run over four scale-row indices for a + // fixed logical column. A 64-row CTA owns two such rows; pack them as a + // uint16 store only when the pair stays inside the same 4-row swizzle group. + const int32_t row_pairs = valid_rows / 2; + const int32_t total_pairs = row_pairs * valid_cols; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_pairs; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = (i / valid_cols) * 2; + int32_t col = i % valid_cols; + const int32_t abs_col = block_coords.x + col; + const int32_t abs_row = scale_row_base + row; + e8m0_t val0 = sColwiseScale[row * CastTraits::blockDIM::N + col]; + e8m0_t val1 = sColwiseScale[(row + 1) * CastTraits::blockDIM::N + col]; + size_t idx = gemm_swizzled_scale_idx(abs_col, abs_row, col_swz_num_tiles_X); + if (((abs_row & 3) != 3) && + ((reinterpret_cast(&scales_colwise[idx]) & (alignof(uint16_t) - 1)) == 0)) { + uint16_t val2 = static_cast(val0) | (static_cast(val1) << 8); + *reinterpret_cast(&scales_colwise[idx]) = val2; + } else { + scales_colwise[idx] = val0; + size_t idx1 = gemm_swizzled_scale_idx(abs_col, abs_row + 1, col_swz_num_tiles_X); + scales_colwise[idx1] = val1; + } + } + // Odd-row tail. + if ((valid_rows & 1) != 0) { + const int32_t row = valid_rows - 1; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < valid_cols; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + const int32_t col = i; + e8m0_t val = sColwiseScale[row * CastTraits::blockDIM::N + col]; + size_t idx = gemm_swizzled_scale_idx(block_coords.x + col, scale_row_base + row, + col_swz_num_tiles_X); + scales_colwise[idx] = val; + } + } + } + + // Cached colwise scale flush (non-swizzled linear layout). + // Rows in scales_colwise are indexed by (block_coords.y / colChunkElems), + // columns are logical input columns; 4 adjacent columns for the same scale + // row live at 4 consecutive gmem bytes, so pack as uint32 stores. + if constexpr (CastTraits::_cache_colwise_scale_in_smem && !CastTraits::_with_swizzled_scales) { + const int32_t scale_row_base = block_coords.y / CastTraits::colChunkElems; + // DIVUP so a partial last block (e.g. rows=993, colChunkElems=32 -> last CTA + // has 1 valid input row) still emits its scale row. Truncating divison here + // drops the last partial scale row, causing the fast path to diverge from + // nvte_swizzle_scaling_factors on non-multiple-of-32-rows inputs. + const int32_t valid_rows = + DIVUP(std::min(block_coords.y + CastTraits::blockDIM::M, rows) - block_coords.y, + static_cast(CastTraits::colChunkElems)); + const int32_t valid_cols = + std::min(block_coords.x + CastTraits::blockDIM::N, cols) - block_coords.x; + + constexpr int32_t cols_per_group = 4; + const int32_t groups_per_row = valid_cols / cols_per_group; + const int32_t total_groups = valid_rows * groups_per_row; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_groups; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / groups_per_row; + int32_t group = i % groups_per_row; + int32_t col = group * cols_per_group; + uint32_t val4 = + *reinterpret_cast(&sColwiseScale[row * CastTraits::blockDIM::N + col]); + size_t idx = + static_cast(scale_row_base + row) * static_cast(scale_stride_colwise) + + static_cast(block_coords.x + col); + *reinterpret_cast(&scales_colwise[idx]) = val4; + } + // Column tail (valid_cols not multiple of 4). + const int32_t remaining_start = groups_per_row * cols_per_group; + if (remaining_start < valid_cols) { + const int32_t remaining_cols = valid_cols - remaining_start; + const int32_t total_remaining = valid_rows * remaining_cols; + for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_remaining; + i += CastTraits::warpLayout::num * THREADS_PER_WARP) { + int32_t row = i / remaining_cols; + int32_t col = remaining_start + (i % remaining_cols); + e8m0_t val = sColwiseScale[row * CastTraits::blockDIM::N + col]; + size_t idx = + static_cast(scale_row_base + row) * static_cast(scale_stride_colwise) + + static_cast(block_coords.x + col); + scales_colwise[idx] = val; } } } @@ -1595,7 +1987,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( ptx::cp_async_bulk_wait_group_read<0>(); #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -} +} // NOLINT(readability/fn_size) } // namespace specialized } // namespace quantize_kernel diff --git a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh index 91c6af26b5..0b980c5086 100644 --- a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh @@ -262,11 +262,9 @@ __global__ void __launch_bounds__(THREADS_NUM) constexpr size_t out_mem_rowwise_scales = 0; extern __shared__ char dynamic_shmem[]; - uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); // Manually align dynamic SHMEM per TMA requirements using padding // __align__(128) Does not guarantee the pointer to be aligned! - uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) & - ~(static_cast(TMA_SHMEM_ALIGNMENT - 1)); + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned IType *in_sh = reinterpret_cast(dshmem); diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index 3a34c76de8..3309b9fa06 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -99,6 +99,31 @@ __launch_bounds__(BLOCK_SIZE) #endif } +template +__global__ void __launch_bounds__(BLOCK_SIZE) + compute_columnwise_amax_kernel(const int num_rows, const int num_cols, + const IType *__restrict__ input, + float *__restrict__ output_columnwise_amax, + const float *__restrict__ noop) { + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + + const int col_idx = blockIdx.x; + if (col_idx >= num_cols) return; + + float thread_max = 0.0f; + for (int row_idx = threadIdx.x; row_idx < num_rows; row_idx += BLOCK_SIZE) { + thread_max = fmaxf(thread_max, fabsf(static_cast(input[row_idx * num_cols + col_idx]))); + } + const float col_amax = + reduce_max(thread_max, threadIdx.x / THREADS_PER_WARP); + + if (threadIdx.x == 0) { + output_columnwise_amax[col_idx] = col_amax; + } +} + template void launch_compute_rowwise_amax(const int num_rows, const int num_cols, const IType *input, float *output_rowwise_amax, cudaStream_t stream, @@ -113,6 +138,20 @@ void launch_compute_rowwise_amax(const int num_rows, const int num_cols, const I NVTE_CHECK_CUDA(cudaGetLastError()); } +template +void launch_compute_columnwise_amax(const int num_rows, const int num_cols, const IType *input, + float *output_columnwise_amax, cudaStream_t stream, + const float *noop = nullptr) { + if (num_rows == 0 || num_cols == 0) return; + + dim3 grid(num_cols); + dim3 block(ROWWISE_AMAX_BLOCK_SIZE); + + compute_columnwise_amax_kernel + <<>>(num_rows, num_cols, input, output_columnwise_amax, noop); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + #endif // FP4_TYPE_SUPPORTED } // namespace rowwise_amax_kernel @@ -155,6 +194,40 @@ inline void compute_rowwise_amax(const Tensor &input, const Tensor *noop, Tensor #endif // FP4_TYPE_SUPPORTED } +inline void compute_columnwise_amax(const Tensor &input, const Tensor *noop, Tensor *output, + cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED + using namespace rowwise_amax_kernel; + + const auto [rows, cols] = input.flat_2d_dims(); + auto *amax_ptr = reinterpret_cast(output->columnwise_amax.dptr); + NVTE_CHECK(amax_ptr != nullptr, "Row-scaled columnwise amax tensor must be allocated."); + NVTE_CHECK(output->columnwise_amax.numel() == cols, "Row-scaled columnwise amax must have ", cols, + " entries, got ", output->columnwise_amax.shape, "."); + + const auto *noop_ptr = reinterpret_cast(noop->data.dptr); + if (input.dtype() == DType::kBFloat16) { + const auto *input_ptr = reinterpret_cast(input.data.dptr); + launch_compute_columnwise_amax<__nv_bfloat16>(static_cast(rows), static_cast(cols), + input_ptr, amax_ptr, stream, noop_ptr); + } else if (input.dtype() == DType::kFloat16) { + const auto *input_ptr = reinterpret_cast(input.data.dptr); + launch_compute_columnwise_amax(static_cast(rows), static_cast(cols), input_ptr, + amax_ptr, stream, noop_ptr); + } else if (input.dtype() == DType::kFloat32) { + const auto *input_ptr = reinterpret_cast(input.data.dptr); + launch_compute_columnwise_amax(static_cast(rows), static_cast(cols), input_ptr, + amax_ptr, stream, noop_ptr); + } else { + NVTE_ERROR( + "Unsupported input dtype for row-scaled NVFP4 quantization. " + "Expected BFloat16, Float16, or Float32."); + } +#else + NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); +#endif // FP4_TYPE_SUPPORTED +} + namespace quantize_transpose_kernel { using namespace quantization_and_transposition_SF; @@ -330,11 +403,9 @@ __global__ void __launch_bounds__(THREADS_NUM) constexpr size_t out_mem_rowwise_scales = 0; extern __shared__ char dynamic_shmem[]; - uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); // Manually align dynamic SHMEM per TMA requirements using padding // __align__(128) Does not guarantee the pointer to be aligned! - uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) & - ~(static_cast(TMA_SHMEM_ALIGNMENT - 1)); + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned IType *in_sh = reinterpret_cast(dshmem); @@ -872,11 +943,9 @@ __global__ void __launch_bounds__(THREADS_NUM) constexpr size_t out_mem_rowwise_scales = 0; extern __shared__ char dynamic_shmem[]; - uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); // Manually align dynamic SHMEM per TMA requirements using padding // __align__(128) Does not guarantee the pointer to be aligned! - uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) & - ~(static_cast(TMA_SHMEM_ALIGNMENT - 1)); + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned IType *in_sh = reinterpret_cast(dshmem); diff --git a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh index 8f37229fd5..ad21486368 100644 --- a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh +++ b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh @@ -184,14 +184,12 @@ compute_nvfp4_scaling_coefficient(const nvfp4_scale_t S_dec_block, const f return static_cast(scale_rcp); } -template -__device__ __forceinline__ void colwise_scaling(const IType *__restrict__ sIn_ptr, - fp4e2m1x2 *__restrict__ sOut_tr_ptr, - nvfp4_scale_t *__restrict__ sSFcolwise_ptr, - const float S_enc_colwise, const int stage_Y, - const int stage_X, const int buff_in, - const int buff_out_tr, RNG_t &rng, - uint4 &random_uint4, int &rnd_idx) { +template +__device__ __forceinline__ void colwise_scaling( + const IType *__restrict__ sIn_ptr, fp4e2m1x2 *__restrict__ sOut_tr_ptr, + nvfp4_scale_t *__restrict__ sSFcolwise_ptr, const float S_enc_colwise, const int stage_Y, + const int stage_X, const int buff_in, const int buff_out_tr, const float *amax_colwise_ptr, + const size_t col_offset, const size_t cols, RNG_t &rng, uint4 &random_uint4, int &rnd_idx) { using scaling_coeff_type = typename SCALING_COEFFICIENT_TYPE::type; const auto &sIn2x = *reinterpret_cast(sIn_ptr); @@ -231,13 +229,21 @@ __device__ __forceinline__ void colwise_scaling(const IType *__restrict__ sIn_pt static_cast(__habs(thread_amax_2x.y))}; #pragma unroll for (int w = 0; w < 2; ++w) { - const nvfp4_scale_t S_dec_b_fp8 = compute_decoding_scaling_factor(block_amax[w], S_enc_colwise); + float S_enc_colwise_block = S_enc_colwise; + if constexpr (ROW_SCALED_NVFP4) { + const size_t col_idx = col_offset + stage_X * TILE_DIM_X + thread_offset_X_colwise + w; + S_enc_colwise_block = + col_idx < cols ? core::compute_global_encode_scaling_factor_FP4(amax_colwise_ptr[col_idx]) + : 1.0f; + } + const nvfp4_scale_t S_dec_b_fp8 = + compute_decoding_scaling_factor(block_amax[w], S_enc_colwise_block); // Store scaling factors to SMEM buffer (R2S) sSFcolwise[scale_tr_offset_Y + w][scale_tr_offset_X] = S_dec_b_fp8; const scaling_coeff_type SFcoefficient = - compute_nvfp4_scaling_coefficient(S_dec_b_fp8, S_enc_colwise); + compute_nvfp4_scaling_coefficient(S_dec_b_fp8, S_enc_colwise_block); // Scale elements __align__(8) uint32_t rOut[SCALE_DIM / 8]; @@ -432,7 +438,7 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D : core::compute_global_encode_scaling_factor_FP4(*amax_rowwise_ptr); const float S_enc_colwise = - (amax_colwise_ptr == nullptr) + (amax_colwise_ptr == nullptr || ROW_SCALED_NVFP4) ? S_enc_rowwise : core::compute_global_encode_scaling_factor_FP4(*amax_colwise_ptr); @@ -587,9 +593,9 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D amax_rowwise_ptr, block_offset_Y, rows, rng, random_uint4, rnd_idx); if constexpr (RETURN_TRANSPOSE) { - colwise_scaling( + colwise_scaling( sIn_ptr, sOut_tr_ptr, sSFcolwise_ptr, S_enc_colwise, stage_Y, stage_X, buff_in, - buff_out_tr, rng, random_uint4, rnd_idx); + buff_out_tr, amax_colwise_ptr, block_offset_X, cols, rng, random_uint4, rnd_idx); } // Wait for shared memory writes to be visible to TMA engine @@ -708,14 +714,14 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop, NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated"); NVTE_CHECK(!row_scaled_nvfp4 || output->amax.dptr != nullptr, "Row-scaled NVFP4 quantization requires rowwise amax."); - NVTE_CHECK(!row_scaled_nvfp4 || !output->has_columnwise_data(), - "Row-scaled NVFP4 quantization does not produce columnwise output."); if (return_transpose) { NVTE_CHECK(is_fp4_dtype(output->columnwise_data.dtype), "Transposed output must have FP4 type."); NVTE_CHECK(output->columnwise_scale_inv.dptr != nullptr, "Transposed scaling tensor must be allocated"); + NVTE_CHECK(!row_scaled_nvfp4 || output->columnwise_amax.dptr != nullptr, + "Row-scaled NVFP4 transpose quantization requires columnwise amax."); } const auto [rows, cols] = input.flat_2d_dims(); diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index babe757da0..a603cfdbef 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -663,6 +663,7 @@ struct QuantizationConfig { bool use_fast_math = false; NVTENVFP44Over6Mode nvfp4_4over6_mode = kNVTENVFP44Over6Disabled; bool nvfp4_4over6_err_use_fast_math = false; + bool mxfp8_2d_quantization = false; #ifdef __HIP_PLATFORM_AMD__ bool mxfp4_use_hadamard = false; #endif @@ -677,7 +678,8 @@ struct QuantizationConfig { sizeof(uint8_t), // stochastic_rounding sizeof(uint8_t), // use_fast_math sizeof(uint8_t), // nvfp4_4over6_mode - sizeof(uint8_t) // nvfp4_4over6_err_use_fast_math + sizeof(uint8_t), // nvfp4_4over6_err_use_fast_math + sizeof(uint8_t) // mxfp8_2d_quantization #ifdef __HIP_PLATFORM_AMD__ , sizeof(uint8_t) // mxfp4_use_hadamard @@ -1271,6 +1273,26 @@ inline bool is_aligned_ptr(const void *ptr, size_t alignment) { return reinterpret_cast(ptr) % alignment == 0; } +/*! \brief Align a shared-memory base pointer up to `align` bytes. + * + * The result is derived from `p` by pointer arithmetic on purpose, without losing its + * identity as a pointer in between, so the address is never rounded through an integer + * -- in which case the compiler would lose the link back to the `extern __shared__` + * object, and ptxas could no longer prove the address lives in the shared window and + * would fall back to generic address-space accesses (`LD.E`/`ST.E`) instead of + * `LDS`/`STS`. + * + * `align` must be a power of two. + */ +__device__ __forceinline__ char *align_up(char *p, uintptr_t align) { + const uintptr_t misalign = reinterpret_cast(p) & (align - 1); + // If p is not aligned, (align - misalign) & (align - 1) is the number of bytes to fill the gap between p and + // the next aligned address. + // If p is aligned, misalign is 0 and (align - misalign) is align itself, so we use & (align - 1) + // to make it 0 and return p itself. + return p + ((align - misalign) & (align - 1)); +} + inline bool is_aligned_tensor_data(const Tensor &t, size_t alignment) { return is_aligned_ptr(static_cast(t.data.dptr), alignment); } diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index a82ec1c98d..b726784d03 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -47,27 +47,52 @@ ncclDataType_t te_dtype_to_nccl_dtype(NVTEDType dtype) { return ncclFloat8e4m3; case kNVTEFloat8E5M2: return ncclFloat8e5m2; + case kNVTEFloat8E8M0: + return ncclUint8; default: NVTE_ERROR("Unsupported NVTEDType for NCCL dtype conversion: ", static_cast(dtype)); } return ncclFloat32; // unreachable } -// shape_out is caller-owned; desc.sizes aliases shape_out.data and must -// outlive the NCCL EP call. +// Which part of a TE tensor an NCCL descriptor points at: the data payload, or +// (for block-scaled tensors) the rowwise scale-inverse that rides alongside it. +enum class DescSource { kData, kScaleInv }; + +// Build an NCCL descriptor for a TE tensor's data (kData) or its rowwise +// scale-inverse (kScaleInv). shape_out is caller-owned; desc.sizes aliases +// shape_out.data and must outlive the NCCL EP call. Uses the matching window +// field (win.window / win.scale_window) when set, else the raw pointer. inline ncclEpTensor_t make_nccl_ep_tensor(const NVTETensor t, NVTEShape& shape_out, - const NVTECommWindow& win = {}) { - shape_out = nvte_tensor_shape(t); + const NVTECommWindow& win = {}, + DescSource source = DescSource::kData) { ncclEpTensor_t desc = NCCL_EP_TENSOR_INIT; + void* raw_ptr = nullptr; + ncclWindow_t win_hdl = nullptr; + uint64_t win_offset = 0; + if (source == DescSource::kData) { + shape_out = nvte_tensor_shape(t); + desc.datatype = te_dtype_to_nccl_dtype(nvte_tensor_type(t)); + raw_ptr = nvte_tensor_data(t); + win_hdl = win.window; + win_offset = win.offset; + } else { + const SimpleTensor& si = convertNVTETensorCheck(t)->scale_inv; + shape_out = nvte_make_shape(si.shape.data(), si.shape.size()); + desc.datatype = te_dtype_to_nccl_dtype(static_cast(si.dtype)); + raw_ptr = si.dptr; + win_hdl = win.scale_window; + win_offset = win.scale_offset; + } desc.ndim = shape_out.ndim; desc.sizes = shape_out.data; - desc.datatype = te_dtype_to_nccl_dtype(nvte_tensor_type(t)); - if (win.window != nullptr) { - desc.win_hdl = win.window; - desc.win_offset = win.offset; + if (win_hdl != nullptr) { + desc.win_hdl = win_hdl; + desc.win_offset = win_offset; } else { - desc.data = nvte_tensor_data(t); - NVTE_CHECK(desc.data != nullptr, "tensor data must not be null"); + desc.data = raw_ptr; + NVTE_CHECK(desc.data != nullptr || nvte_tensor_numel(t) == 0, + "non-empty tensor data must not be null"); } return desc; } @@ -94,8 +119,15 @@ void EPBackend::validate_config(const NVTEEpGroupConfig& config) { NVTE_CHECK(config.num_experts > 0, "num_experts must be positive, got ", config.num_experts); NVTE_CHECK(config.max_tokens_per_rank > 0, "max_tokens_per_rank must be positive, got ", config.max_tokens_per_rank); - NVTE_CHECK(config.max_recv_tokens_per_rank > 0, "max_recv_tokens_per_rank must be positive, got ", + // 0 selects eager mode (NCCL_EP_AUTO); any explicit budget must be positive. + NVTE_CHECK(config.max_recv_tokens_per_rank >= 0, + "max_recv_tokens_per_rank must be non-negative, got ", config.max_recv_tokens_per_rank); + NVTE_CHECK(!(config.zero_copy && config.max_recv_tokens_per_rank == 0), + "zero-copy and eager (max_recv_tokens_per_rank = 0) modes are mutually exclusive"); + NVTE_CHECK(!(config.drop_on_overflow && config.max_recv_tokens_per_rank == 0), + "drop_on_overflow (overflow drop) is not supported in eager mode " + "(max_recv_tokens_per_rank = 0)"); NVTE_CHECK(config.hidden_dim > 0, "hidden_dim must be positive, got ", config.hidden_dim); NVTE_CHECK(config.max_token_dtype >= 0 && config.max_token_dtype < kNVTENumTypes, "max_token_dtype out of range, got ", static_cast(config.max_token_dtype)); @@ -210,9 +242,13 @@ void EPBackend::init(ncclComm_t ep_comm, NVTEEpGroupConfig group_config) { cfg.max_num_sms = group_config.num_comm_sms > 0 ? static_cast(group_config.num_comm_sms) : NCCL_EP_AUTO; - // Must be > 0; NCCL EP errors out on 0. + // 0 = NCCL_EP_AUTO, which enables eager mode (recv buffers sized per routing). cfg.max_recv_tokens_per_rank = static_cast(group_config.max_recv_tokens_per_rank); cfg.zero_copy = group_config.zero_copy ? NCCL_EP_ZERO_COPY_ON : NCCL_EP_ZERO_COPY_OFF; + // Per-token top-k; NCCL EP sizes internal buffers from it in eager mode. + cfg.num_topk = static_cast(group_config.num_topk); + cfg.overflow_policy = + group_config.drop_on_overflow ? NCCL_EP_OVERFLOW_DROP : NCCL_EP_OVERFLOW_AUTO; NVTE_CHECK_NCCL(ncclEpCreateGroup(&ep_group_, ep_comm, &cfg)); @@ -320,10 +356,8 @@ size_t EPBackend::handle_mem_size(NVTEEpLayerConfig layer_cfg) { } void EPBackend::prepare(void* handle_mem, const NVTETensor topk_idx, - NVTETensor recv_tokens_per_expert, - NVTETensor /*total_recv_tokens_per_rank*/, NVTEEpLayerConfig layer_cfg, - cudaStream_t stream) { - // total_recv_tokens_per_rank is a reserved placeholder; not yet populated. + NVTETensor recv_tokens_per_expert, NVTETensor total_recv_tokens_per_rank, + NVTEEpLayerConfig layer_cfg, cudaStream_t stream) { NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); NVTE_CHECK(layer_cfg.top_k > 0, "top_k must be > 0, got ", layer_cfg.top_k); NVTE_CHECK(nvte_tensor_shape(topk_idx).ndim == 2, "topk_idx must be 2D [T, top_k]"); @@ -331,16 +365,24 @@ void EPBackend::prepare(void* handle_mem, const NVTETensor topk_idx, NVTEShape topk_idx_shape; ncclEpTensor_t nccl_topk_idx = make_nccl_ep_tensor(topk_idx, topk_idx_shape); - // ncclEpUpdateHandle writes per-expert counts via expert_counters. + // ncclEpUpdateHandle writes per-expert counts via expert_counters and, when + // provided, the scalar padded recv-slot total via recv_total_counter. NVTEShape recv_tokens_per_expert_shape; ncclEpTensor_t recv_tokens_per_expert_desc; if (recv_tokens_per_expert != nullptr) { recv_tokens_per_expert_desc = make_nccl_ep_tensor(recv_tokens_per_expert, recv_tokens_per_expert_shape); } + NVTEShape total_recv_shape; + ncclEpTensor_t total_recv_desc; + if (total_recv_tokens_per_rank != nullptr) { + total_recv_desc = make_nccl_ep_tensor(total_recv_tokens_per_rank, total_recv_shape); + } ncclEpLayoutInfo_t layout_info = NCCL_EP_LAYOUT_INFO_INIT; layout_info.expert_counters = (recv_tokens_per_expert != nullptr) ? &recv_tokens_per_expert_desc : nullptr; + layout_info.recv_total_counter = + (total_recv_tokens_per_rank != nullptr) ? &total_recv_desc : nullptr; std::lock_guard lock(mutex_); NVTE_CHECK(initialized_, "EPBackend not initialized"); @@ -394,16 +436,44 @@ void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTE make_nccl_ep_tensor(recv_topk_weights, recv_topk_weights_shape, recv_topk_weights_win); } + // Block-scaled (e.g. MXFP8): route the per-token scale-inverse alongside the + // data. High-precision (bf16/fp16/fp32) and per-tensor FP8 payloads carry the + // default delayed scaling mode and skip this. Keys on is_block_scaling so + // NVFP4 can reuse this path later. + const NVTEScalingMode tokens_scaling_mode = nvte_tensor_scaling_mode(tokens); + const bool is_scaled = is_block_scaling(tokens_scaling_mode); + NVTEShape scales_in_shape, scales_out_shape; + ncclEpTensor_t nccl_scales_in = NCCL_EP_TENSOR_INIT, nccl_scales_out = NCCL_EP_TENSOR_INIT; + if (is_scaled) { + NVTE_CHECK(is_mxfp8_scaling(tokens_scaling_mode), + "EP dispatch supports MXFP8 block scaling only; got scaling mode ", + static_cast(tokens_scaling_mode)); + NVTE_CHECK(nvte_tensor_scaling_mode(recv_tokens) == tokens_scaling_mode, + "recv_tokens scaling mode must match tokens scaling mode"); + nccl_scales_in = + make_nccl_ep_tensor(tokens, scales_in_shape, tokens_win, DescSource::kScaleInv); + nccl_scales_out = + make_nccl_ep_tensor(recv_tokens, scales_out_shape, recv_tokens_win, DescSource::kScaleInv); + } else { + NVTE_CHECK(!is_fp8_dtype(static_cast(tok_dtype)), + "EP dispatch of FP8 tokens requires a block scaling mode (e.g. MXFP8); " + "per-tensor (delayed) FP8 scaling is not supported"); + } + ncclEpDispatchInputs_t in_struct = NCCL_EP_DISPATCH_INPUTS_INIT; in_struct.tokens = &nccl_tokens_in; in_struct.topk_weights = is_forward ? &nccl_topk_weights_in : nullptr; + in_struct.scales = is_scaled ? &nccl_scales_in : nullptr; ncclEpDispatchOutputs_t out_struct = NCCL_EP_DISPATCH_OUTPUTS_INIT; out_struct.tokens = &nccl_tokens_out; out_struct.topk_weights = is_forward ? &nccl_topk_weights_out : nullptr; + out_struct.scales = is_scaled ? &nccl_scales_out : nullptr; ncclEpDispatchConfig_t dispatch_cfg = NCCL_EP_DISPATCH_CONFIG_INIT; dispatch_cfg.pass_direction = is_forward ? NCCL_EP_FWD_PASS : NCCL_EP_BWD_PASS; + // Block-scaled payloads forward the per-token scale-inverse; select the matching recipe. + dispatch_cfg.quant_recipe = is_scaled ? NCCL_EP_DISP_QUANT_FWD : NCCL_EP_DISP_QUANT_NONE; std::lock_guard lock(mutex_); NVTE_CHECK(initialized_, "EPBackend not initialized"); diff --git a/transformer_engine/common/fused_attn/context_parallel.cu b/transformer_engine/common/fused_attn/context_parallel.cu index 0f7b820bbd..a6cf76ffe1 100644 --- a/transformer_engine/common/fused_attn/context_parallel.cu +++ b/transformer_engine/common/fused_attn/context_parallel.cu @@ -42,6 +42,12 @@ struct CopyFunctor { } }; +struct ZeroFunctor { + __forceinline__ __device__ static void run(void *token, void *token_per_step, int idx) { + reinterpret_cast(token)[idx] = make_float4(0.f, 0.f, 0.f, 0.f); + } +}; + template struct AddFunctor { __forceinline__ __device__ static void run(dtype *token, dtype *token_per_step, int idx) { @@ -357,24 +363,27 @@ __global__ void thd_grad_correction_kernel(dtype *grad, dtype *grad_per_step, in for (int token_id = group_id; token_id < num_total_tokens; token_id += num_groups) { int seq_id = binary_search(token_id, cu_seqlens_s, batch + 1); - int token_offset; - bool is_first_half; if constexpr (functor_idx < 2) { - token_offset = cu_seqlens_s[seq_id + functor_idx]; - is_first_half = (functor_idx == 0); + dtype *first_half_token = + &grad[(token_id + cu_seqlens_s[seq_id]) * static_cast(hidden_size)]; + dtype *second_half_token = + &grad[(token_id + cu_seqlens_s[seq_id + 1]) * static_cast(hidden_size)]; + dtype *token_per_step = &grad_per_step[token_id * static_cast(hidden_size)]; + for (int idx = lane_id; idx < num_inner_loops; idx += group_size) { + Functor_0::run(first_half_token, token_per_step, idx); + Functor_1::run(second_half_token, token_per_step, idx); + } } else { - token_offset = 0; int len = cu_seqlens_s[seq_id + 1] - cu_seqlens_s[seq_id]; - is_first_half = (token_id - cu_seqlens_s[seq_id]) < (len / 2); - } - - dtype *token = &grad[(token_id + token_offset) * static_cast(hidden_size)]; - dtype *token_per_step = &grad_per_step[token_id * static_cast(hidden_size)]; - for (int idx = lane_id; idx < num_inner_loops; idx += group_size) { - if (is_first_half) { - Functor_0::run(token, token_per_step, idx); - } else { - Functor_1::run(token, token_per_step, idx); + bool is_first_half = (token_id - cu_seqlens_s[seq_id]) < (len / 2); + dtype *token = &grad[token_id * static_cast(hidden_size)]; + dtype *token_per_step = &grad_per_step[token_id * static_cast(hidden_size)]; + for (int idx = lane_id; idx < num_inner_loops; idx += group_size) { + if (is_first_half) { + Functor_0::run(token, token_per_step, idx); + } else { + Functor_1::run(token, token_per_step, idx); + } } } } @@ -707,6 +716,12 @@ static void thd_grad_dispatcher(Tensor grad, const Tensor &grad_per_step, const } else if (first_half == "none" && second_half == "copy") { thd_grad_correction_helper(grad, grad_per_step, cu_seqlens, stream); + } else if (first_half == "copy" && second_half == "zero") { + thd_grad_correction_helper(grad, grad_per_step, cu_seqlens, + stream); + } else if (first_half == "zero" && second_half == "copy") { + thd_grad_correction_helper(grad, grad_per_step, cu_seqlens, + stream); } else if (first_half == "add" && second_half == "copy") { thd_grad_correction_helper, CopyFunctor, 2>(grad, grad_per_step, cu_seqlens, stream); @@ -722,6 +737,19 @@ void thd_grad_correction(Tensor grad, const Tensor &grad_per_step, const Tensor const std::string &first_half, const std::string &second_half, cudaStream_t stream) { using namespace transformer_engine; + if (grad.dtype() == DType::kByte) { + if (first_half == "copy" && second_half == "zero") { + thd_grad_correction_helper(grad, grad_per_step, cu_seqlens, + stream); + } else if (first_half == "zero" && second_half == "copy") { + thd_grad_correction_helper(grad, grad_per_step, cu_seqlens, + stream); + } else { + NVTE_ERROR( + "FP8 gradients stored as raw encoded bytes require copy/zero or zero/copy correction\n"); + } + return; + } TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( grad.dtype(), dtype, thd_grad_dispatcher(grad, grad_per_step, cu_seqlens, first_half, second_half, diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index fc21771297..544c02aceb 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -240,6 +240,8 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + const bool is_thd_layout = + q_format == NVTE_QKV_Format::NVTE_THD || kv_format == NVTE_QKV_Format::NVTE_THD; NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); auto cudnn_runtime_version = cudnnGetVersion(); @@ -272,7 +274,9 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( (attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK || attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)) || + attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || + (sm_arch_ >= 100 && + attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK))) || // 9.21: d_qk=192, d_v=128 (cudnn_runtime_version >= 92100 && sm_arch_ >= 100 && head_dim_qk <= 192 && head_dim_v <= 128 && head_dim_qk % 16 == 0 && head_dim_v % 16 == 0 && @@ -281,13 +285,18 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK))) && // pre-9.21: {bshd, sbhd}, {vanilla} // 9.21+: {bshd, sbhd, bhsd}, {vanilla, off-by-one, learnable} + // 9.23+: {thd}; sm90 fwd only, sm100+ fwd/bwd ((cudnn_runtime_version < 92100 && (qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD) && softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX) || (cudnn_runtime_version >= 92100 && (qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD || - qkv_format == NVTE_QKV_Format::NVTE_BHSD))) && - !requires_64bit_ragged_offset && + qkv_format == NVTE_QKV_Format::NVTE_BHSD)) || + ((cudnn_runtime_version >= 92300 && (sm_arch_ >= 100 || (sm_arch_ >= 90 && !is_training))) && + qkv_format == NVTE_QKV_Format::NVTE_THD && supported_ragged_offset_size && + (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || + attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || + attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK))) && // 9.10.0: known bugs with SDPA FP8 (cudnn_runtime_version != 91000) && !return_max_logit) { backend = NVTE_Fused_Attn_Backend::NVTE_FP8; @@ -327,7 +336,23 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK))) || // 9.11: d_qk = 192, d_v = 128 + Blackwell + bprop + non-paged (head_dim_qk == 192 && head_dim_v == 128 && is_training && sm_arch_ >= 100 && - cudnn_runtime_version >= 91100)) && + cudnn_runtime_version >= 91100) || + // 9.23: d_qk = d_v = 256 + SM10x (cuDNN FE 1.24 / BE 9.23+) + bprop + non-paged. + // THD layouts require cuDNN FE 1.26 / BE 9.25+ for execution-plan support. + (head_dim_qk == 256 && head_dim_v == 256 && is_training && sm_arch_ >= 100 && + sm_arch_ < 110 && cudnn_runtime_version >= (is_thd_layout ? 92500 : 92300) && + layout_group != NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD && + // The FE forces this path onto the deterministic bprop algorithm, which on + // Blackwell rejects dBias, dropout, and ALiBi (and supports vanilla softmax only). + bias_type == NVTE_Bias_Type::NVTE_NO_BIAS && dropout == 0.0 && + softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX && + // Non-causal D=256 supports only full-window attention; SWA is allowed only for causal masks. + ((window_size_left == -1 && window_size_right == -1) || + ((attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || + attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || + attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK || + attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK) && + (window_size_right == -1 || window_size_right == 0))))) && // 9.11+ bug: 128 < d_qk <= 256, 128 < d_v <= 256 + Hopper + bprop + MLA // Conditional to temporarily use blanket cudnn_runtime_version >= 9.11 until fixed (!((cudnn_runtime_version >= 91100) && is_training && sm_arch_ == 90 && @@ -627,12 +652,13 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso input_cu_seqlens_kv, input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, input_page_table_k, input_page_table_v, input_rng_state, wkspace, stream, handle); } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_FP8) { - fused_attn_fp8_fwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, is_training, + fused_attn_fp8_fwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, t_q, t_kv, is_training, attn_scale, dropout, qkv_layout, o_format, qkv_scale_inv_format, bias_type, attn_mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, input_Q, input_K, input_V, input_SoftmaxOffset, input_output_S, output_O, Aux_CTX_Tensors, input_cu_seqlens_q, - input_cu_seqlens_kv, input_rng_state, wkspace, stream, handle); + input_cu_seqlens_kv, input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, + input_rng_state, wkspace, stream, handle); } else { NVTE_ERROR("Invalid combination of data type and sequence length for fused attention. \n"); } @@ -729,14 +755,15 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso if (input_dO->scaling_mode == NVTE_MXFP8_1D_SCALING) { input_dO_f16 = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); } - fused_attn_fp8_bwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, attn_scale, dropout, - qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, + fused_attn_fp8_bwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, t_q, t_kv, attn_scale, + dropout, qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, do_scale_inv_format, bias_type, attn_mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, deterministic, input_Q, input_K, input_V, input_O, input_dO, input_dO_f16, input_M, input_S, input_SoftmaxOffset, input_output_dP, output_dQ, output_dK, output_dV, output_dSoftmaxOffset, input_cu_seqlens_q, input_cu_seqlens_kv, - input_rng_state, wkspace, stream, handle); + input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, input_rng_state, + wkspace, stream, handle); } else { NVTE_ERROR("Invalid combination of data type and sequence length for fused attention. \n"); } diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 000af41aee..2ef0ac3393 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -6,6 +6,7 @@ #include "../common.h" #include "../cudnn_utils.h" +#include "../util/cuda_runtime.h" #include "../util/system.h" #include "fused_attn_fp8.h" #include "utils.h" @@ -15,17 +16,20 @@ namespace fused_attn { using namespace transformer_engine; +constexpr size_t kFP8THDRaggedCudnnVersion = 92300; + // fused attention FWD FP8 with FE 1.0+ void fused_attn_fp8_fwd_impl( int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v, - bool is_training, float scaling_factor, float dropout_probability, NVTE_QKV_Layout qkv_layout, - NVTE_QKV_Format o_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, void* devPtrQ, void* devPtrK, void* devPtrV, - void* devPtrSoftmaxOffset, void* devPtrM, void* devPtrO, void* devPtrDescaleQ, - void* devPtrDescaleK, void* devPtrDescaleV, void* devPtrDescaleS, void* devPtrScaleS, - void* devPtrScaleO, void* devPtrAmaxO, void* devPtrAmaxS, void* devPtrcuSeqlensQ, - void* devPtrcuSeqlensKV, void* devPtrDropoutSeed, void* devPtrDropoutOffset, + int64_t max_b, int64_t max_t_q, int64_t max_t_kv, bool is_training, float scaling_factor, + float dropout_probability, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, + NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, + int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, void* devPtrQ, + void* devPtrK, void* devPtrV, void* devPtrSoftmaxOffset, void* devPtrM, void* devPtrO, + void* devPtrDescaleQ, void* devPtrDescaleK, void* devPtrDescaleV, void* devPtrDescaleS, + void* devPtrScaleS, void* devPtrScaleO, void* devPtrAmaxO, void* devPtrAmaxS, + void* devPtrcuSeqlensQ, void* devPtrcuSeqlensKV, void* devPtrSeqOffsetsQ, + void* devPtrSeqOffsetsKV, void* devPtrDropoutSeed, void* devPtrDropoutOffset, cudnn_frontend::DataType_t qkv_tensor_type, cudnn_frontend::DataType_t o_tensor_type, NVTEScalingMode scaling_mode, NVTE_QKV_Format qkv_scale_inv_format, void* workspace, size_t* workspace_size, cudaStream_t stream, cudnnHandle_t handle) { @@ -34,9 +38,11 @@ void fused_attn_fp8_fwd_impl( bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); bool is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); + (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) || + (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)); bool is_padding = ((mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); + (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) || + (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)); bool is_dropout = (is_training && dropout_probability != 0.0f); bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); auto bias_b = b; @@ -60,11 +66,21 @@ void fused_attn_fp8_fwd_impl( NVTE_CHECK(!is_mxfp8 || cudnn_runtime_version >= 92100, "MXFP8 fused attention requires cuDNN 9.21.0 or later!"); + NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); + NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + bool is_ragged_q = (q_format == NVTE_QKV_Format::NVTE_THD); + bool is_ragged_kv = (kv_format == NVTE_QKV_Format::NVTE_THD); + const int device_id = cuda::current_device(); + const int sm_arch_ = cuda::sm_arch(device_id); + bool use_ragged_stats = + is_ragged_q && cudnn_runtime_version >= kFP8THDRaggedCudnnVersion && sm_arch_ != 120; + + NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); + const DType ragged_offset_type = DType::kInt64; + // Newer versions of cuDNN SDPA can accept sequence lengths directly as a cumulative - // tensor. Take advantage of this if possible to avoid 1 extra kernel call. (Unlike - // the F16 path, the FP8 path has no THD/ragged-offset support, so only the - // cu_seqlens_to_actual_seqlens conversion applies here. Also note that the - // needed versions of cuDNN backend and frontend are higher than for F16.) + // tensor. Take advantage of this if possible to avoid the actual-seqlen conversion; + // THD inputs still use their separate ragged-offset tensors. const bool use_cu_seqlens_directly = // Frontend 1.26 supports fp8+cu_seqlens (for the C++ API). // Note: For the Python API, 1.27 is required. @@ -78,6 +94,20 @@ void fused_attn_fp8_fwd_impl( // (which doesn't support cu_seqlens). Remove this restriction when possible. !is_dropout; + int64_t actual_b = b; + if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= kFP8THDRaggedCudnnVersion) { + NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); + if (sm_arch_ != 120) { + // cuDNN reads the user's [actual_b+1] cu_seqlens buffers directly, so a quantized + // batch dimension would read out of bounds on the direct path. + if (!use_cu_seqlens_directly) { + b = max_b; + } + s_q = is_ragged_q ? max_t_q : s_q; + s_kv = is_ragged_kv ? max_t_kv : s_kv; + } + } + try { FADescriptor_v1 descriptor{b, h, @@ -139,6 +169,11 @@ void fused_attn_fp8_fwd_impl( std::shared_ptr, // softmax_offset std::shared_ptr, // seq_q std::shared_ptr, // seq_kv + std::shared_ptr, // offset_q + std::shared_ptr, // offset_k + std::shared_ptr, // offset_v + std::shared_ptr, // offset_o + std::shared_ptr, // offset_stats std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset @@ -164,6 +199,8 @@ void fused_attn_fp8_fwd_impl( std::shared_ptr descale_q, descale_k, descale_v; std::shared_ptr descale_s, scale_s, scale_o; std::shared_ptr bias, softmax_offset, seq_q, seq_kv; + std::shared_ptr offset_q, offset_k, offset_v, offset_o, + offset_stats; std::shared_ptr dropout_seed, dropout_offset; // Q, K, V, attn_scale @@ -175,6 +212,14 @@ void fused_attn_fp8_fwd_impl( .set_dim({b, h, s_q, d_qk}) .set_stride(q_strides) .set_data_type(qkv_tensor_type)); + if (is_ragged_q) { + offset_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_q") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + Q->set_ragged_offset(offset_q); + } K = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("K") .set_dim({b, hg, s_kv, d_qk}) @@ -185,6 +230,20 @@ void fused_attn_fp8_fwd_impl( .set_dim({b, hg, s_kv, d_v}) .set_stride(v_strides) .set_data_type(qkv_tensor_type)); + if (is_ragged_kv) { + offset_k = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_k") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + offset_v = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_v") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + K->set_ragged_offset(offset_k); + V->set_ragged_offset(offset_v); + } attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("attn_scale") .set_dim({1, 1, 1, 1}) @@ -362,15 +421,33 @@ void fused_attn_fp8_fwd_impl( .set_dim({b, h, s_q, d_v}) .set_stride(o_strides) .set_data_type(o_tensor_type); + if (is_ragged_q) { + offset_o = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_o") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + O->set_ragged_offset(offset_o); + } amax_o->set_output(!is_mxfp8) .set_dim({1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) .set_data_type(fe::DataType_t::FLOAT); - Stats->set_output(true) - .set_data_type(fe::DataType_t::FLOAT) - .set_dim({b, h, s_q, 1}) - .set_stride({h * s_q, s_q, 1, 1}); + if (use_ragged_stats) { + offset_stats = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_stats") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + } + Stats->set_output(true).set_data_type(fe::DataType_t::FLOAT).set_dim({b, h, s_q, 1}); + if (use_ragged_stats) { + Stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); + } else { + Stats->set_stride({h * s_q, s_q, 1, 1}); + } std::tuple, // Q std::shared_ptr, // K @@ -396,6 +473,12 @@ void fused_attn_fp8_fwd_impl( is_softmax_offset ? std::make_tuple(softmax_offset) : std::make_tuple(nullptr); auto padding_tuple = is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); + auto offset_q_tuple = is_ragged_q ? std::make_tuple(offset_q) : std::make_tuple(nullptr); + auto offset_kv_tuple = + is_ragged_kv ? std::make_tuple(offset_k, offset_v) : std::make_tuple(nullptr, nullptr); + auto offset_o_tuple = is_ragged_q ? std::make_tuple(offset_o) : std::make_tuple(nullptr); + auto offset_s_tuple = + use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) : std::make_tuple(nullptr, nullptr); @@ -406,24 +489,37 @@ void fused_attn_fp8_fwd_impl( NVTE_CHECK_CUDNN_FE(mha_graph->build_plans(handle)); auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, - softmax_offset_tuple, padding_tuple, dropout_tuple); + softmax_offset_tuple, padding_tuple, offset_q_tuple, offset_kv_tuple, + offset_o_tuple, offset_s_tuple, dropout_tuple); cache.insert({descriptor, return_tuple}); return return_tuple; }; auto [mha_graph, Q, K, V, descale_q, descale_k, descale_v, descale_s, scale_s, scale_o, - attn_scale, O, amax_s, amax_o, Stats, bias, softmax_offset, seq_q, seq_kv, dropout_seed, - dropout_offset] = get_graph(sdpa_fp8_fprop_cache, descriptor); - - auto plan_workspace_size = mha_graph->get_workspace_size(); + attn_scale, O, amax_s, amax_o, Stats, bias, softmax_offset, seq_q, seq_kv, offset_q, + offset_k, offset_v, offset_o, offset_stats, dropout_seed, dropout_offset] = + get_graph(sdpa_fp8_fprop_cache, descriptor); + + auto plan_workspace_size = alignTo<16>(mha_graph->get_workspace_size()); + const size_t num_bytes_per_seqlen = alignTo<16>(b * sizeof(int32_t)); + const size_t actual_seqlen_workspace_size = + (is_padding && !use_cu_seqlens_directly) ? 2 * num_bytes_per_seqlen : 0; + const size_t num_bytes_per_ragged_offset = + alignTo<16>(((b + 1) * typeToNumBits(ragged_offset_type)) / 8); + size_t seqlen_offsets_workspace_size = 0; + if (is_ragged_q || is_ragged_kv) { + size_t count = 2 * (static_cast(is_ragged_q) + static_cast(is_ragged_kv)); + if (use_ragged_stats) { + seqlen_offsets_workspace_size = (count + 1) * num_bytes_per_ragged_offset; + } else { + seqlen_offsets_workspace_size = count * num_bytes_per_ragged_offset; + } + } - // Exit to request upper level API to allocate memory if needed. - // When passing cu_seqlens* directly to cuDNN SDPA, no conversion workspace is - // needed: cuDNN consumes the user's cu_seqlens buffers as-is. - size_t actual_seqlen_workspace_size = use_cu_seqlens_directly ? 0 : 2 * b * sizeof(int32_t); if (workspace == nullptr) { - *workspace_size = plan_workspace_size + actual_seqlen_workspace_size; + *workspace_size = + plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; return; } @@ -465,9 +561,9 @@ void fused_attn_fp8_fwd_impl( constexpr size_t nthreads_per_block = 128; const size_t grid = (b + nthreads_per_block - 1) / nthreads_per_block; void* devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; - void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + b * sizeof(int32_t); + void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + num_bytes_per_seqlen; cu_seqlens_to_actual_seqlens<<>>( - b, b, static_cast(devPtrcuSeqlensQ), // TODO(pass max_b) + actual_b, b, static_cast(devPtrcuSeqlensQ), static_cast(devPtrcuSeqlensKV), static_cast(devActualSeqlenQ), static_cast(devActualSeqlenKV)); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -476,6 +572,49 @@ void fused_attn_fp8_fwd_impl( } } + if (is_ragged_q || is_ragged_kv) { + constexpr size_t nthreads_per_block = 128; + const size_t grid = (b + nthreads_per_block) / nthreads_per_block; + void* devOffsets = + static_cast(workspace) + plan_workspace_size + actual_seqlen_workspace_size; + void* devOffsetsQ = nullptr; + void* devOffsetsO = nullptr; + if (is_ragged_q) { + devOffsetsQ = devOffsets; + devOffsetsO = static_cast(devOffsetsQ) + num_bytes_per_ragged_offset; + } + void* devOffsetsK = nullptr; + void* devOffsetsV = nullptr; + if (is_ragged_kv) { + devOffsetsK = static_cast(devOffsets) + + static_cast(is_ragged_q) * 2 * num_bytes_per_ragged_offset; + devOffsetsV = static_cast(devOffsetsK) + num_bytes_per_ragged_offset; + } + void* devOffsetsS = nullptr; + if (use_ragged_stats) { + devOffsetsS = static_cast(devOffsets) + + (static_cast(is_ragged_q) + static_cast(is_ragged_kv)) * 2 * + num_bytes_per_ragged_offset; + } + const RaggedOffsetMultipliers offset_mults(layout_group, h, hg, d_qk, d_v); + cu_seqlens_padded_to_offsets<<>>( + offset_mults, actual_b, b, static_cast(devPtrSeqOffsetsQ), + static_cast(devPtrSeqOffsetsKV), ragged_offset_type, devOffsetsQ, devOffsetsK, + devOffsetsV, devOffsetsO, devOffsetsS); + NVTE_CHECK_CUDA(cudaGetLastError()); + if (is_ragged_q) { + variant_pack[offset_q] = devOffsetsQ; + variant_pack[offset_o] = devOffsetsO; + } + if (is_ragged_kv) { + variant_pack[offset_k] = devOffsetsK; + variant_pack[offset_v] = devOffsetsV; + } + if (use_ragged_stats) { + variant_pack[offset_stats] = devOffsetsS; + } + } + if (is_dropout) { variant_pack[dropout_seed] = devPtrDropoutSeed; variant_pack[dropout_offset] = devPtrDropoutOffset; @@ -489,37 +628,40 @@ void fused_attn_fp8_fwd_impl( } catch (cudnn_frontend::cudnnException& e) { NVTE_ERROR(e.what()); } -} +} // NOLINT(readability/fn_size) // fused attention BWD FP8 with FE 1.0+ void fused_attn_fp8_bwd_impl( int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v, - float scaling_factor, float dropout_probability, NVTE_QKV_Layout qkv_layout, - NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, - int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, - bool deterministic, void* devPtrQ, void* devPtrK, void* devPtrV, void* devPtrM, void* devPtrO, - void* devPtrdO, void* devPtrSoftmaxOffset, void* devPtrdQ, void* devPtrdK, void* devPtrdV, + int64_t max_b, int64_t max_t_q, int64_t max_t_kv, float scaling_factor, + float dropout_probability, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, + NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, NVTE_Bias_Type bias_type, + NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, + int64_t window_size_right, bool bottom_right_diagonal, bool deterministic, void* devPtrQ, + void* devPtrK, void* devPtrV, void* devPtrM, void* devPtrO, void* devPtrdO, + void* devPtrSoftmaxOffset, void* devPtrdQ, void* devPtrdK, void* devPtrdV, void* devPtrdSoftmaxOffset, void* devPtrDescaleQ, void* devPtrDescaleK, void* devPtrDescaleV, void* devPtrDescaleO, void* devPtrDescaledO, void* devPtrDescaleS, void* devPtrDescaledP, void* devPtrScaleS, void* devPtrScaledP, void* devPtrScaledQ, void* devPtrScaledK, void* devPtrScaledV, void* devPtrAmaxdP, void* devPtrAmaxdQ, void* devPtrAmaxdK, void* devPtrAmaxdV, void* devPtrQ_t, void* devPtrK_t, void* devPtrdO_f16, void* devPtrdO_t, void* devPtrDescaleQ_t, void* devPtrDescaleK_t, void* devPtrDescaledO_t, void* devPtrcuSeqlensQ, - void* devPtrcuSeqlensKV, void* devPtrDropoutSeed, void* devPtrDropoutOffset, - cudnn_frontend::DataType_t qkv_tensor_type, cudnn_frontend::DataType_t o_tensor_type, - cudnn_frontend::DataType_t do_tensor_type, cudnn_frontend::DataType_t dqkv_tensor_type, - NVTEScalingMode scaling_mode, NVTE_QKV_Format qkv_scale_inv_format, - NVTE_QKV_Format do_scale_inv_format, void* workspace, size_t* workspace_size, - cudaStream_t stream, cudnnHandle_t handle) { + void* devPtrcuSeqlensKV, void* devPtrSeqOffsetsQ, void* devPtrSeqOffsetsKV, + void* devPtrDropoutSeed, void* devPtrDropoutOffset, cudnn_frontend::DataType_t qkv_tensor_type, + cudnn_frontend::DataType_t o_tensor_type, cudnn_frontend::DataType_t do_tensor_type, + cudnn_frontend::DataType_t dqkv_tensor_type, NVTEScalingMode scaling_mode, + NVTE_QKV_Format qkv_scale_inv_format, NVTE_QKV_Format do_scale_inv_format, void* workspace, + size_t* workspace_size, cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; const auto cudnn_runtime_version = cudnnGetVersion(); bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); bool is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); + (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) || + (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)); bool is_padding = ((mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); + (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) || + (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)); bool is_dropout = (dropout_probability != 0.0f); bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); auto bias_b = b; @@ -543,6 +685,28 @@ void fused_attn_fp8_bwd_impl( NVTE_CHECK(!is_mxfp8 || cudnn_runtime_version >= 92100, "MXFP8 fused attention requires cuDNN 9.21.0 or later!"); + NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); + NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + bool is_ragged_q = (q_format == NVTE_QKV_Format::NVTE_THD); + bool is_ragged_kv = (kv_format == NVTE_QKV_Format::NVTE_THD); + const int device_id = cuda::current_device(); + const int sm_arch_ = cuda::sm_arch(device_id); + bool use_ragged_stats = + is_ragged_q && cudnn_runtime_version >= kFP8THDRaggedCudnnVersion && sm_arch_ != 120; + + NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); + const DType ragged_offset_type = DType::kInt64; + + int64_t actual_b = b; + if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= kFP8THDRaggedCudnnVersion) { + NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); + if (sm_arch_ != 120) { + b = max_b; + s_q = is_ragged_q ? max_t_q : s_q; + s_kv = is_ragged_kv ? max_t_kv : s_kv; + } + } + bool is_O_in_F16 = (o_tensor_type == cudnn_frontend::DataType_t::HALF || o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); @@ -628,6 +792,11 @@ void fused_attn_fp8_bwd_impl( std::shared_ptr, // d_softmax_offset std::shared_ptr, // seq_q std::shared_ptr, // seq_kv + std::shared_ptr, // offset_q + std::shared_ptr, // offset_k + std::shared_ptr, // offset_v + std::shared_ptr, // offset_o + std::shared_ptr, // offset_stats std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset @@ -660,6 +829,8 @@ void fused_attn_fp8_bwd_impl( std::shared_ptr scale_dQ, scale_dK, scale_dV; std::shared_ptr bias, dBias, softmax_offset, d_softmax_offset; std::shared_ptr seq_q, seq_kv; + std::shared_ptr offset_q, offset_k, offset_v, offset_o, + offset_stats; std::shared_ptr dropout_seed, dropout_offset; // Q, K, V, O, dO, stats, attn_scale @@ -673,6 +844,19 @@ void fused_attn_fp8_bwd_impl( .set_dim({b, h, s_q, d_qk}) .set_stride(q_strides) .set_data_type(qkv_tensor_type)); + if (is_ragged_q) { + offset_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_q") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + offset_o = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_o") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + Q->set_ragged_offset(offset_q); + } K = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("K") .set_dim({b, hg, s_kv, d_qk}) @@ -683,21 +867,53 @@ void fused_attn_fp8_bwd_impl( .set_dim({b, hg, s_kv, d_v}) .set_stride(v_strides) .set_data_type(qkv_tensor_type)); + if (is_ragged_kv) { + offset_k = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_k") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + offset_v = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_v") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + K->set_ragged_offset(offset_k); + V->set_ragged_offset(offset_v); + } O = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("O") .set_dim({b, h, s_q, d_v}) .set_stride(o_strides) .set_data_type(o_tensor_type)); + if (is_ragged_q) { + O->set_ragged_offset(offset_o); + } dO = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("dO") .set_dim({b, h, s_q, d_v}) .set_stride(dO_strides) .set_data_type(do_tensor_type)); + if (is_ragged_q) { + dO->set_ragged_offset(offset_o); + } + if (use_ragged_stats) { + offset_stats = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_stats") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + } Stats = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("Stats") .set_dim({b, h, s_q, 1}) - .set_stride({h * s_q, s_q, 1, 1}) .set_data_type(fe::DataType_t::FLOAT)); + if (use_ragged_stats) { + Stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); + } else { + Stats->set_stride({h * s_q, s_q, 1, 1}); + } attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("attn_scale") .set_dim({1, 1, 1, 1}) @@ -947,6 +1163,9 @@ void fused_attn_fp8_bwd_impl( .set_dim({b, h, s_q, d_qk}) .set_stride(dq_strides) .set_data_type(dqkv_tensor_type); + if (is_ragged_q) { + dQ->set_ragged_offset(offset_q); + } dK->set_output(true) .set_dim({b, hg, s_kv, d_qk}) .set_stride(dk_strides) @@ -955,6 +1174,10 @@ void fused_attn_fp8_bwd_impl( .set_dim({b, hg, s_kv, d_v}) .set_stride(dv_strides) .set_data_type(dqkv_tensor_type); + if (is_ragged_kv) { + dK->set_ragged_offset(offset_k); + dV->set_ragged_offset(offset_v); + } amax_dQ->set_output(!is_mxfp8) .set_dim({1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) @@ -1013,6 +1236,12 @@ void fused_attn_fp8_bwd_impl( : std::make_tuple(nullptr, nullptr); auto padding_tuple = is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); + auto offset_q_tuple = is_ragged_q ? std::make_tuple(offset_q) : std::make_tuple(nullptr); + auto offset_kv_tuple = + is_ragged_kv ? std::make_tuple(offset_k, offset_v) : std::make_tuple(nullptr, nullptr); + auto offset_o_tuple = is_ragged_q ? std::make_tuple(offset_o) : std::make_tuple(nullptr); + auto offset_s_tuple = + use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) : std::make_tuple(nullptr, nullptr); @@ -1024,7 +1253,8 @@ void fused_attn_fp8_bwd_impl( auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, mxfp8_tensors_tuple, - bias_tuple, softmax_offset_tuple, padding_tuple, dropout_tuple); + bias_tuple, softmax_offset_tuple, padding_tuple, offset_q_tuple, + offset_kv_tuple, offset_o_tuple, offset_s_tuple, dropout_tuple); cache.insert({descriptor, return_tuple}); return return_tuple; @@ -1033,14 +1263,27 @@ void fused_attn_fp8_bwd_impl( descale_dO, descale_s, descale_dP, scale_s, scale_dQ, scale_dK, scale_dV, scale_dP, dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP, Q_t, K_t, dO_f16, dO_t, descale_q_t, descale_k_t, descale_dO_t, bias, dBias, softmax_offset, d_softmax_offset, seq_q, seq_kv, - dropout_seed, dropout_offset] = get_graph(sdpa_fp8_bprop_cache, descriptor); - - auto plan_workspace_size = mha_graph->get_workspace_size(); + offset_q, offset_k, offset_v, offset_o, offset_stats, dropout_seed, dropout_offset] = + get_graph(sdpa_fp8_bprop_cache, descriptor); + + auto plan_workspace_size = alignTo<16>(mha_graph->get_workspace_size()); + const size_t num_bytes_per_seqlen = alignTo<16>(b * sizeof(int32_t)); + const size_t actual_seqlen_workspace_size = is_padding ? 2 * num_bytes_per_seqlen : 0; + const size_t num_bytes_per_ragged_offset = + alignTo<16>(((b + 1) * typeToNumBits(ragged_offset_type)) / 8); + size_t seqlen_offsets_workspace_size = 0; + if (is_ragged_q || is_ragged_kv) { + size_t count = 2 * (static_cast(is_ragged_q) + static_cast(is_ragged_kv)); + if (use_ragged_stats) { + seqlen_offsets_workspace_size = (count + 1) * num_bytes_per_ragged_offset; + } else { + seqlen_offsets_workspace_size = count * num_bytes_per_ragged_offset; + } + } - // Exit to request upper level API to allocate memory if needed - size_t actual_seqlen_workspace_size = 2 * b * sizeof(int32_t); if (workspace == nullptr) { - *workspace_size = plan_workspace_size + actual_seqlen_workspace_size; + *workspace_size = + plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; return; } @@ -1106,9 +1349,9 @@ void fused_attn_fp8_bwd_impl( constexpr size_t nthreads_per_block = 128; const size_t grid = (b + nthreads_per_block - 1) / nthreads_per_block; void* devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; - void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + b * sizeof(int32_t); + void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + num_bytes_per_seqlen; cu_seqlens_to_actual_seqlens<<>>( - b, b, static_cast(devPtrcuSeqlensQ), // TODO(pass max_b) + actual_b, b, static_cast(devPtrcuSeqlensQ), static_cast(devPtrcuSeqlensKV), static_cast(devActualSeqlenQ), static_cast(devActualSeqlenKV)); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -1116,6 +1359,49 @@ void fused_attn_fp8_bwd_impl( variant_pack[seq_kv] = devActualSeqlenKV; } + if (is_ragged_q || is_ragged_kv) { + constexpr size_t nthreads_per_block = 128; + const size_t grid = (b + nthreads_per_block) / nthreads_per_block; + void* devOffsets = + static_cast(workspace) + plan_workspace_size + actual_seqlen_workspace_size; + void* devOffsetsQ = nullptr; + void* devOffsetsO = nullptr; + if (is_ragged_q) { + devOffsetsQ = devOffsets; + devOffsetsO = static_cast(devOffsetsQ) + num_bytes_per_ragged_offset; + } + void* devOffsetsK = nullptr; + void* devOffsetsV = nullptr; + if (is_ragged_kv) { + devOffsetsK = static_cast(devOffsets) + + static_cast(is_ragged_q) * 2 * num_bytes_per_ragged_offset; + devOffsetsV = static_cast(devOffsetsK) + num_bytes_per_ragged_offset; + } + void* devOffsetsS = nullptr; + if (use_ragged_stats) { + devOffsetsS = static_cast(devOffsets) + + (static_cast(is_ragged_q) + static_cast(is_ragged_kv)) * 2 * + num_bytes_per_ragged_offset; + } + const RaggedOffsetMultipliers offset_mults(layout_group, h, hg, d_qk, d_v); + cu_seqlens_padded_to_offsets<<>>( + offset_mults, actual_b, b, static_cast(devPtrSeqOffsetsQ), + static_cast(devPtrSeqOffsetsKV), ragged_offset_type, devOffsetsQ, devOffsetsK, + devOffsetsV, devOffsetsO, devOffsetsS); + NVTE_CHECK_CUDA(cudaGetLastError()); + if (is_ragged_q) { + variant_pack[offset_q] = devOffsetsQ; + variant_pack[offset_o] = devOffsetsO; + } + if (is_ragged_kv) { + variant_pack[offset_k] = devOffsetsK; + variant_pack[offset_v] = devOffsetsV; + } + if (use_ragged_stats) { + variant_pack[offset_stats] = devOffsetsS; + } + } + if (is_dropout) { variant_pack[dropout_seed] = devPtrDropoutSeed; variant_pack[dropout_offset] = devPtrDropoutOffset; @@ -1137,14 +1423,16 @@ void fused_attn_fp8_bwd_impl( // fused attention FWD FP8 with separate Q, K, V void fused_attn_fp8_fwd( size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, - size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, bool is_training, float attn_scale, - float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, - NVTE_QKV_Format qkv_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, size_t window_size_left, size_t window_size_right, - bool bottom_right_diagonal, const Tensor* input_Q, const Tensor* input_K, const Tensor* input_V, + size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, size_t num_tokens_q, + size_t num_tokens_kv, bool is_training, float attn_scale, float p_dropout, + NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format qkv_scale_inv_format, + NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, + size_t window_size_left, size_t window_size_right, bool bottom_right_diagonal, + const Tensor* input_Q, const Tensor* input_K, const Tensor* input_V, const Tensor* input_SoftmaxOffset, Tensor* input_output_S, Tensor* output_O, NVTETensorPack* Aux_CTX_Tensors, const Tensor* cu_seqlens_q, const Tensor* cu_seqlens_kv, - const Tensor* rng_state, Tensor* workspace, cudaStream_t stream, cudnnHandle_t handle) { + const Tensor* cu_seqlens_q_padded, const Tensor* cu_seqlens_kv_padded, const Tensor* rng_state, + Tensor* workspace, cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; void *devPtrQ = nullptr, *devPtrK = nullptr, *devPtrV = nullptr; void *devPtrDescaleQ = nullptr, *devPtrDescaleK = nullptr, *devPtrDescaleV = nullptr; @@ -1171,12 +1459,39 @@ void fused_attn_fp8_fwd( if (softmax_type != NVTE_VANILLA_SOFTMAX) { devPtrSoftmaxOffset = input_SoftmaxOffset->data.dptr; } + NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); + NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + const auto cudnn_runtime_version = cudnnGetVersion(); + const int sm_arch_ = cuda::sm_arch(cuda::current_device()); + + void* devPtrSeqOffsetsQ = cu_seqlens_q_padded->data.dptr; + void* devPtrSeqOffsetsKV = cu_seqlens_kv_padded->data.dptr; + + size_t max_batch_size = 0; + size_t max_tokens_q = 0; + size_t max_tokens_kv = 0; + if (q_format == NVTE_QKV_Format::NVTE_THD || kv_format == NVTE_QKV_Format::NVTE_THD) { + max_batch_size = fused_attn::get_max_batch_size(batch); + } + if (q_format == NVTE_QKV_Format::NVTE_THD) { + max_tokens_q = fused_attn::get_max_tokens(num_tokens_q); + } + if (kv_format == NVTE_QKV_Format::NVTE_THD) { + max_tokens_kv = fused_attn::get_max_tokens(num_tokens_kv); + } + void* devPtrM = nullptr; if (Aux_CTX_Tensors->size == 0) { int i = 0; Tensor* output_M = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_M->data.dptr = nullptr; - output_M->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; + // SM120 uses dense stats in the graph, so its allocation must remain [b, h, s_q, 1]. + if (q_format == NVTE_QKV_Format::NVTE_THD && + cudnn_runtime_version >= fused_attn::kFP8THDRaggedCudnnVersion && sm_arch_ != 120) { + output_M->data.shape = {num_tokens_q, num_attn_heads, 1}; + } else { + output_M->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; + } output_M->data.dtype = DType::kFloat32; Tensor* output_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_rng_state->data.dptr = nullptr; @@ -1218,18 +1533,19 @@ void fused_attn_fp8_fwd( NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); if ((qkv_format == NVTE_QKV_Format::NVTE_BSHD) || (qkv_format == NVTE_QKV_Format::NVTE_SBHD) || - (qkv_format == NVTE_QKV_Format::NVTE_BHSD)) { + (qkv_format == NVTE_QKV_Format::NVTE_BHSD) || (qkv_format == NVTE_QKV_Format::NVTE_THD)) { fused_attn::fused_attn_fp8_fwd_impl( batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, - is_training, attn_scale, p_dropout, qkv_layout, o_format, bias_type, mask_type, - softmax_type, window_size_left, window_size_right, bottom_right_diagonal, devPtrQ, devPtrK, - devPtrV, devPtrSoftmaxOffset, devPtrM, devPtrO, devPtrDescaleQ, devPtrDescaleK, - devPtrDescaleV, devPtrDescaleS, devPtrScaleS, devPtrScaleO, devPtrAmaxO, devPtrAmaxS, - devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrDropoutSeed, devPtrDropoutOffset, - get_cudnn_fe_dtype(QKV_type), get_cudnn_fe_dtype(O_type), input_Q->scaling_mode, - qkv_scale_inv_format, workspace->data.dptr, &workspace_size, stream, handle); + max_batch_size, max_tokens_q, max_tokens_kv, is_training, attn_scale, p_dropout, qkv_layout, + o_format, bias_type, mask_type, softmax_type, window_size_left, window_size_right, + bottom_right_diagonal, devPtrQ, devPtrK, devPtrV, devPtrSoftmaxOffset, devPtrM, devPtrO, + devPtrDescaleQ, devPtrDescaleK, devPtrDescaleV, devPtrDescaleS, devPtrScaleS, devPtrScaleO, + devPtrAmaxO, devPtrAmaxS, devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrSeqOffsetsQ, + devPtrSeqOffsetsKV, devPtrDropoutSeed, devPtrDropoutOffset, get_cudnn_fe_dtype(QKV_type), + get_cudnn_fe_dtype(O_type), input_Q->scaling_mode, qkv_scale_inv_format, + workspace->data.dptr, &workspace_size, stream, handle); } else { - NVTE_ERROR("FP8 fused attention only supports qkv_format=BSHD, SBHD, or BHSD.\n"); + NVTE_ERROR("FP8 fused attention only supports qkv_format=BSHD, SBHD, BHSD, or THD.\n"); } if (workspace_size > 0) { @@ -1247,18 +1563,20 @@ void fused_attn_fp8_fwd( // fused attention BWD FP8 with separate Q, K, V void fused_attn_fp8_bwd( size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, - size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, float attn_scale, float p_dropout, - NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, - NVTE_QKV_Layout dqkv_layout, NVTE_QKV_Format qkv_scale_inv_format, - NVTE_QKV_Format do_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, size_t window_size_left, size_t window_size_right, - bool bottom_right_diagonal, bool deterministic, const Tensor* input_Q, const Tensor* input_K, - const Tensor* input_V, const Tensor* input_O, const Tensor* input_dO, - const Tensor* input_dO_f16, const Tensor* input_M, const Tensor* input_S, - const Tensor* input_SoftmaxOffset, Tensor* input_output_dP, const Tensor* output_dQ, - const Tensor* output_dK, const Tensor* output_dV, Tensor* output_dSoftmaxOffset, - const Tensor* cu_seqlens_q, const Tensor* cu_seqlens_kv, const Tensor* rng_state, - Tensor* workspace, cudaStream_t stream, cudnnHandle_t handle) { + size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, size_t num_tokens_q, + size_t num_tokens_kv, float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, + NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, + NVTE_QKV_Format qkv_scale_inv_format, NVTE_QKV_Format do_scale_inv_format, + NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, + size_t window_size_left, size_t window_size_right, bool bottom_right_diagonal, + bool deterministic, const Tensor* input_Q, const Tensor* input_K, const Tensor* input_V, + const Tensor* input_O, const Tensor* input_dO, const Tensor* input_dO_f16, + const Tensor* input_M, const Tensor* input_S, const Tensor* input_SoftmaxOffset, + Tensor* input_output_dP, const Tensor* output_dQ, const Tensor* output_dK, + const Tensor* output_dV, Tensor* output_dSoftmaxOffset, const Tensor* cu_seqlens_q, + const Tensor* cu_seqlens_kv, const Tensor* cu_seqlens_q_padded, + const Tensor* cu_seqlens_kv_padded, const Tensor* rng_state, Tensor* workspace, + cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; void* devPtrQ = input_Q->data.dptr; void* devPtrK = input_K->data.dptr; @@ -1323,6 +1641,25 @@ void fused_attn_fp8_bwd( devPtrScaledV = output_dV->scale.dptr; } + NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); + NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + + void* devPtrSeqOffsetsQ = cu_seqlens_q_padded->data.dptr; + void* devPtrSeqOffsetsKV = cu_seqlens_kv_padded->data.dptr; + + size_t max_batch_size = 0; + size_t max_tokens_q = 0; + size_t max_tokens_kv = 0; + if (q_format == NVTE_QKV_Format::NVTE_THD || kv_format == NVTE_QKV_Format::NVTE_THD) { + max_batch_size = fused_attn::get_max_batch_size(batch); + } + if (q_format == NVTE_QKV_Format::NVTE_THD) { + max_tokens_q = fused_attn::get_max_tokens(num_tokens_q); + } + if (kv_format == NVTE_QKV_Format::NVTE_THD) { + max_tokens_kv = fused_attn::get_max_tokens(num_tokens_kv); + } + void* devPtrcuSeqlensQ = reinterpret_cast(reinterpret_cast(cu_seqlens_q->data.dptr)); void* devPtrcuSeqlensKV = @@ -1339,23 +1676,24 @@ void fused_attn_fp8_bwd( NVTE_QKV_Format dqkv_format = nvte_get_qkv_format(dqkv_layout); if ((dqkv_format == NVTE_QKV_Format::NVTE_BSHD) || (dqkv_format == NVTE_QKV_Format::NVTE_SBHD) || - (dqkv_format == NVTE_QKV_Format::NVTE_BHSD)) { + (dqkv_format == NVTE_QKV_Format::NVTE_BHSD) || (dqkv_format == NVTE_QKV_Format::NVTE_THD)) { fused_attn::fused_attn_fp8_bwd_impl( batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, - attn_scale, p_dropout, qkv_layout, o_format, do_format, dqkv_layout, bias_type, mask_type, - softmax_type, window_size_left, window_size_right, bottom_right_diagonal, deterministic, - devPtrQ, devPtrK, devPtrV, devPtrM, devPtrO, devPtrdO, devPtrSoftmaxOffset, devPtrdQ, - devPtrdK, devPtrdV, devPtrdSoftmaxOffset, devPtrDescaleQ, devPtrDescaleK, devPtrDescaleV, - devPtrDescaleO, devPtrDescaledO, devPtrDescaleS, devPtrDescaledP, devPtrScaleS, - devPtrScaledP, devPtrScaledQ, devPtrScaledK, devPtrScaledV, devPtrAmaxdP, devPtrAmaxdQ, - devPtrAmaxdK, devPtrAmaxdV, devPtrQ_t, devPtrK_t, devPtrdO_f16, devPtrdO_t, - devPtrDescaleQ_t, devPtrDescaleK_t, devPtrDescaledO_t, devPtrcuSeqlensQ, devPtrcuSeqlensKV, + max_batch_size, max_tokens_q, max_tokens_kv, attn_scale, p_dropout, qkv_layout, o_format, + do_format, dqkv_layout, bias_type, mask_type, softmax_type, window_size_left, + window_size_right, bottom_right_diagonal, deterministic, devPtrQ, devPtrK, devPtrV, devPtrM, + devPtrO, devPtrdO, devPtrSoftmaxOffset, devPtrdQ, devPtrdK, devPtrdV, devPtrdSoftmaxOffset, + devPtrDescaleQ, devPtrDescaleK, devPtrDescaleV, devPtrDescaleO, devPtrDescaledO, + devPtrDescaleS, devPtrDescaledP, devPtrScaleS, devPtrScaledP, devPtrScaledQ, devPtrScaledK, + devPtrScaledV, devPtrAmaxdP, devPtrAmaxdQ, devPtrAmaxdK, devPtrAmaxdV, devPtrQ_t, devPtrK_t, + devPtrdO_f16, devPtrdO_t, devPtrDescaleQ_t, devPtrDescaleK_t, devPtrDescaledO_t, + devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, devPtrDropoutSeed, devPtrDropoutOffset, get_cudnn_fe_dtype(QKV_type), get_cudnn_fe_dtype(O_type), get_cudnn_fe_dtype(dO_type), get_cudnn_fe_dtype(dQKV_type), input_dO->scaling_mode, qkv_scale_inv_format, do_scale_inv_format, workspace->data.dptr, &workspace_size, stream, handle); } else { - NVTE_ERROR("FP8 fused attention only supports dqkv_format=BSHD, SBHD, or BHSD.\n"); + NVTE_ERROR("FP8 fused attention only supports dqkv_format=BSHD, SBHD, BHSD, or THD.\n"); } if (workspace_size > 0) { diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h index b9660128ca..5b3b2fff14 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -15,28 +15,32 @@ namespace transformer_engine { // fused attention FWD FP8 with separate Q, K, V void fused_attn_fp8_fwd( size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, - size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, bool is_training, float attn_scale, - float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, - NVTE_QKV_Format qkv_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, size_t window_size_left, size_t window_size_right, - bool bottom_right_diagonal, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, + size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, size_t num_tokens_q, + size_t num_tokens_kv, bool is_training, float attn_scale, float p_dropout, + NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format qkv_scale_inv_format, + NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, + size_t window_size_left, size_t window_size_right, bool bottom_right_diagonal, + const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, const Tensor *input_SoftmaxOffset, Tensor *input_output_S, Tensor *output_O, NVTETensorPack *Aux_CTX_Tensors, const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, - const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); + const Tensor *cu_seqlens_q_padded, const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, + Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); // fused attention BWD FP8 with separate Q, K, V void fused_attn_fp8_bwd( size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, - size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, float attn_scale, float p_dropout, - NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, - NVTE_QKV_Layout dqkv_layout, NVTE_QKV_Format qkv_scale_inv_format, - NVTE_QKV_Format do_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, size_t window_size_left, size_t window_size_right, - bool bottom_right_diagonal, bool deterministic, const Tensor *input_Q, const Tensor *input_K, - const Tensor *input_V, const Tensor *input_O, const Tensor *input_dO, - const Tensor *input_dO_f16, const Tensor *input_M, const Tensor *input_S, - const Tensor *input_SoftmaxOffset, Tensor *input_output_dP, const Tensor *output_dQ, - const Tensor *output_dK, const Tensor *output_dV, Tensor *output_dSoftmaxOffset, - const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, const Tensor *rng_state, - Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); + size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, size_t num_tokens_q, + size_t num_tokens_kv, float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, + NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, + NVTE_QKV_Format qkv_scale_inv_format, NVTE_QKV_Format do_scale_inv_format, + NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, + size_t window_size_left, size_t window_size_right, bool bottom_right_diagonal, + bool deterministic, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, + const Tensor *input_O, const Tensor *input_dO, const Tensor *input_dO_f16, + const Tensor *input_M, const Tensor *input_S, const Tensor *input_SoftmaxOffset, + Tensor *input_output_dP, const Tensor *output_dQ, const Tensor *output_dK, + const Tensor *output_dV, Tensor *output_dSoftmaxOffset, const Tensor *cu_seqlens_q, + const Tensor *cu_seqlens_kv, const Tensor *cu_seqlens_q_padded, + const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, Tensor *workspace, + cudaStream_t stream, cudnnHandle_t handle); } // namespace transformer_engine diff --git a/transformer_engine/common/fused_rope/fused_rope.cu b/transformer_engine/common/fused_rope/fused_rope.cu index 27dc11ab43..c7b762e973 100644 --- a/transformer_engine/common/fused_rope/fused_rope.cu +++ b/transformer_engine/common/fused_rope/fused_rope.cu @@ -14,6 +14,27 @@ namespace transformer_engine { +// Returns the largest sequence index `b` such that +// `cu_seqlens[b] / cp_size <= t_id`. Used by the linear-grid THD kernels to +// locate the sequence span that owns local packed token `t_id`. Uses the same +// integer-division semantics as the existing THD kernels so that the +// per-sequence boundaries agree exactly. +__device__ __forceinline__ int fused_rope_thd_find_seq_id(const int *cu_seqlens, const int nseq, + const int t_id, const int cp_size) { + int lo = 0; + int hi = nseq; + while (lo + 1 < hi) { + int mid = (lo + hi) >> 1; + int mid_start = cu_seqlens[mid] / cp_size; + if (mid_start <= t_id) { + lo = mid; + } else { + hi = mid; + } + } + return lo; +} + template __device__ void fused_rope_block_forward(const scalar_t *src, const float *freqs, scalar_t *dst, const bool interleaved, const int s_id, @@ -215,6 +236,119 @@ __global__ void fused_rope_backward_kernel( offset_block_dst, h, d, d2, stride_h, stride_d, o_stride_h, o_stride_d); } +// THD linear-grid forward kernel. Each block handles exactly one packed local +// token row. The block locates its owning sequence via binary search over the +// divided cumulative sequence boundaries, then defers to the same +// `fused_rope_block_forward` device function as the original kernel. +template +__global__ void fused_rope_thd_linear_grid_forward_kernel( + const scalar_t *src, const int *cu_seqlens, const float *freqs, const int *start_positions, + scalar_t *dst, const bool interleaved, const int cp_size, const int cp_rank, const int nseq, + const int h, const int d, const int d2, const int stride_t, const int stride_h, + const int stride_d, const int o_stride_t, const int o_stride_h, const int o_stride_d) { + int t_id = blockIdx.x; + __shared__ int valid_token; + __shared__ int seq_id; + if (threadIdx.x == 0 && threadIdx.y == 0) { + valid_token = t_id < cu_seqlens[nseq] / cp_size; + seq_id = valid_token ? fused_rope_thd_find_seq_id(cu_seqlens, nseq, t_id, cp_size) : 0; + } + __syncthreads(); + if (!valid_token) return; + + int b_id = seq_id; + int start = cu_seqlens[b_id] / cp_size; + int end = cu_seqlens[b_id + 1] / cp_size; + int s_id = t_id - start; + int cur_seqlens = end - start; + + int offset_block = t_id * stride_t; + int offset_block_dst = t_id * o_stride_t; + + int begin_offset = (start_positions == nullptr) ? 0 : start_positions[b_id]; + int s_id_for_freqs = s_id + begin_offset; + + if (cp_size > 1) { + assert(cur_seqlens % 2 == 0); + if (s_id < cur_seqlens / 2) { + s_id_for_freqs += cp_rank * cur_seqlens / 2; + } else { + s_id_for_freqs += cur_seqlens * cp_size - (cp_rank + 1) * cur_seqlens / 2 - cur_seqlens / 2; + } + } + + fused_rope_block_forward(src, freqs, dst, interleaved, s_id_for_freqs, offset_block, + offset_block_dst, h, d, d2, stride_h, stride_d, o_stride_h, o_stride_d); +} + +// THD linear-grid backward kernel. Mirrors the forward variant and dispatches +// to `fused_rope_block_backward`. +template +__global__ void fused_rope_thd_linear_grid_backward_kernel( + const scalar_t *src, const int *cu_seqlens, const float *freqs, const int *start_positions, + scalar_t *dst, const bool interleaved, const int cp_size, const int cp_rank, const int nseq, + const int h, const int d, const int d2, const int stride_t, const int stride_h, + const int stride_d, const int o_stride_t, const int o_stride_h, const int o_stride_d) { + int t_id = blockIdx.x; + __shared__ int valid_token; + __shared__ int seq_id; + if (threadIdx.x == 0 && threadIdx.y == 0) { + valid_token = t_id < cu_seqlens[nseq] / cp_size; + seq_id = valid_token ? fused_rope_thd_find_seq_id(cu_seqlens, nseq, t_id, cp_size) : 0; + } + __syncthreads(); + if (!valid_token) return; + + int b_id = seq_id; + int start = cu_seqlens[b_id] / cp_size; + int end = cu_seqlens[b_id + 1] / cp_size; + int s_id = t_id - start; + int cur_seqlens = end - start; + + int offset_block = t_id * stride_t; + int offset_block_dst = t_id * o_stride_t; + + int begin_offset = (start_positions == nullptr) ? 0 : start_positions[b_id]; + int s_id_for_freqs = s_id + begin_offset; + + if (cp_size > 1) { + assert(cur_seqlens % 2 == 0); + if (s_id < cur_seqlens / 2) { + s_id_for_freqs += cp_rank * cur_seqlens / 2; + } else { + s_id_for_freqs += cur_seqlens * cp_size - (cp_rank + 1) * cur_seqlens / 2 - cur_seqlens / 2; + } + } + + fused_rope_block_backward(src, freqs, dst, interleaved, s_id_for_freqs, offset_block, + offset_block_dst, h, d, d2, stride_h, stride_d, o_stride_h, o_stride_d); +} + +// Host-side dispatcher. Selects the THD linear-grid path when it would +// eliminate a meaningful number of dead blocks. +constexpr size_t kTHDLinearGridOverlaunchThreshold = 2; +constexpr size_t kTHDLinearGridMinWastedBlocks = 65536; + +inline bool use_fused_rope_thd_linear_grid_launch(const NVTE_QKV_Format qkv_format, + const size_t legacy_grid_blocks, + const size_t linear_grid_blocks, + const int cp_size) { + if (qkv_format != NVTE_QKV_Format::NVTE_THD) return false; + if (linear_grid_blocks == 0) return false; + + // Heuristic: use the linear-grid path when the original THD grid, + // `dim3(s, b)`, has both enough relative and absolute block wastage to + // amortize one sequence lookup per useful token. The CP factor keeps the + // ratio gate conservative because rows in the input shrink with context + // parallelism while the original `s * b` launch space does not. + const bool meets_ratio = + legacy_grid_blocks > + kTHDLinearGridOverlaunchThreshold * static_cast(cp_size) * linear_grid_blocks; + const size_t wasted_blocks = + legacy_grid_blocks > linear_grid_blocks ? legacy_grid_blocks - linear_grid_blocks : 0; + return meets_ratio && wasted_blocks >= kTHDLinearGridMinWastedBlocks; +} + template __device__ void fused_qkv_rope_block_forward(const scalar_t *src, const float *freqs, scalar_t *out, const bool interleaved, const int s_id, @@ -467,9 +601,8 @@ void fused_rope_forward_launcher(const scalar_t *input, const int *cu_seqlens, c const int cp_size, const int cp_rank, const int s, const int b, const int h, const int d, const int d2, const int stride_s_or_t, const int stride_b, const int stride_h, const int stride_d, - cudaStream_t stream) { + const int64_t total_tokens_in_input, cudaStream_t stream) { int warps_per_block = h < 16 ? 4 : 8; - dim3 blocks(s, b); dim3 threads(THREADS_PER_WARP, warps_per_block); const int shared_mem_size = 2 * d2 * sizeof(float); // cos, sin int o_stride_s_or_t, o_stride_b; @@ -487,6 +620,19 @@ void fused_rope_forward_launcher(const scalar_t *input, const int *cu_seqlens, c const int o_stride_h = d; const int o_stride_d = 1; + const size_t linear_grid_blocks = static_cast(total_tokens_in_input); + const size_t legacy_grid_blocks = static_cast(s) * static_cast(b); + if (use_fused_rope_thd_linear_grid_launch(qkv_format, legacy_grid_blocks, linear_grid_blocks, + cp_size)) { + dim3 blocks(static_cast(linear_grid_blocks)); + fused_rope_thd_linear_grid_forward_kernel<<>>( + input, cu_seqlens, freqs, start_positions, output, interleaved, cp_size, cp_rank, b, h, d, + d2, stride_s_or_t, stride_h, stride_d, o_stride_s_or_t, o_stride_h, o_stride_d); + NVTE_CHECK_CUDA(cudaGetLastError()); + return; + } + + dim3 blocks(s, b); fused_rope_forward_kernel<<>>( input, cu_seqlens, freqs, start_positions, output, interleaved, cp_size, cp_rank, s, h, d, d2, stride_s_or_t, stride_b, stride_h, stride_d, o_stride_s_or_t, o_stride_b, o_stride_h, @@ -501,9 +647,9 @@ void fused_rope_backward_launcher(const scalar_t *output_grads, const int *cu_se const bool interleaved, const int cp_size, const int cp_rank, const int s, const int b, const int h, const int d, const int d2, const int stride_s_or_t, const int stride_b, const int stride_h, - const int stride_d, cudaStream_t stream) { + const int stride_d, const int64_t total_tokens_in_input, + cudaStream_t stream) { int warps_per_block = h < 16 ? 4 : 8; - dim3 blocks(s, b); dim3 threads(THREADS_PER_WARP, warps_per_block); const int shared_mem_size = 2 * d2 * sizeof(float); // cos, sin int o_stride_s_or_t, o_stride_b; @@ -521,6 +667,20 @@ void fused_rope_backward_launcher(const scalar_t *output_grads, const int *cu_se const int o_stride_h = d; const int o_stride_d = 1; + const size_t linear_grid_blocks = static_cast(total_tokens_in_input); + const size_t legacy_grid_blocks = static_cast(s) * static_cast(b); + if (use_fused_rope_thd_linear_grid_launch(qkv_format, legacy_grid_blocks, linear_grid_blocks, + cp_size)) { + dim3 blocks(static_cast(linear_grid_blocks)); + fused_rope_thd_linear_grid_backward_kernel<<>>( + output_grads, cu_seqlens, freqs, start_positions, input_grads, interleaved, cp_size, + cp_rank, b, h, d, d2, stride_s_or_t, stride_h, stride_d, o_stride_s_or_t, o_stride_h, + o_stride_d); + NVTE_CHECK_CUDA(cudaGetLastError()); + return; + } + + dim3 blocks(s, b); fused_rope_backward_kernel<<>>( output_grads, cu_seqlens, freqs, start_positions, input_grads, interleaved, cp_size, cp_rank, s, h, d, d2, stride_s_or_t, stride_b, stride_h, stride_d, o_stride_s_or_t, o_stride_b, @@ -579,6 +739,12 @@ void fused_rope_forward(const Tensor &input, const Tensor &cu_seqlens, const Ten const int cp_rank, const int s, const int b, const int h, const int d, const int d2, const int stride_s_or_t, const int stride_b, const int stride_h, const int stride_d, cudaStream_t stream) { + // For THD the total packed tokens in the input is the first dimension of the + // tensor. SBHD/BSHD ignore this value. + const int64_t total_tokens_in_input = + (qkv_format == NVTE_QKV_Format::NVTE_THD && !input.data.shape.empty()) + ? static_cast(input.data.shape[0]) + : 0; TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( input.data.dtype, scalar_t, fused_rope_forward_launcher(reinterpret_cast(input.data.dptr), @@ -587,7 +753,7 @@ void fused_rope_forward(const Tensor &input, const Tensor &cu_seqlens, const Ten reinterpret_cast(start_positions.data.dptr), reinterpret_cast(output->data.dptr), qkv_format, interleaved, cp_size, cp_rank, s, b, h, d, d2, stride_s_or_t, - stride_b, stride_h, stride_d, stream);); + stride_b, stride_h, stride_d, total_tokens_in_input, stream);); } void fused_rope_backward(const Tensor &output_grads, const Tensor &cu_seqlens, const Tensor &freqs, @@ -597,6 +763,10 @@ void fused_rope_backward(const Tensor &output_grads, const Tensor &cu_seqlens, c const int h, const int d, const int d2, const int stride_s_or_t, const int stride_b, const int stride_h, const int stride_d, cudaStream_t stream) { + const int64_t total_tokens_in_input = + (qkv_format == NVTE_QKV_Format::NVTE_THD && !output_grads.data.shape.empty()) + ? static_cast(output_grads.data.shape[0]) + : 0; TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( output_grads.data.dtype, scalar_t, fused_rope_backward_launcher(reinterpret_cast(output_grads.data.dptr), @@ -605,7 +775,7 @@ void fused_rope_backward(const Tensor &output_grads, const Tensor &cu_seqlens, c reinterpret_cast(start_positions.data.dptr), reinterpret_cast(input_grads->data.dptr), qkv_format, interleaved, cp_size, cp_rank, s, b, h, d, d2, stride_s_or_t, - stride_b, stride_h, stride_d, stream);); + stride_b, stride_h, stride_d, total_tokens_in_input, stream);); } void fused_qkv_rope_forward(const Tensor &qkv_input, const Tensor &q_freqs, const Tensor &k_freqs, diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu index 412e4aba16..684b0184ff 100644 --- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -1328,7 +1328,8 @@ __global__ void setup_grouped_gemm_kernel( char *a_base, char *b_base, char *c_base, char *d_base, TensorShapeInfo A_meta, TensorShapeInfo B_meta, TensorShapeInfo C_meta, TensorShapeInfo D_meta, size_t a_bits_per_elem, size_t b_bits_per_elem, size_t c_elem_size, size_t d_elem_size, float *alpha_ptr, - float *beta_ptr, bool use_per_group_alpha_beta, + float *beta_ptr, bool use_per_group_alpha_beta, bool a_is_discrete, bool c_is_discrete, + bool d_is_discrete, // Scale inputs: for tensor scaling, pass float* and set mxfp8_base to nullptr // For MXFP8, pass nullptr for tensor_scale and set mxfp8_base float *a_scale_base, float *b_scale_base, bool a_rowwise, bool b_rowwise, @@ -1343,12 +1344,9 @@ __global__ void setup_grouped_gemm_kernel( if (idx >= num_tensors) return; // Get dimensions for this tensor (from array or uniform value) - const bool has_a_multi_tensor = (a_base == nullptr); - const bool has_c_multi_tensor = (c_base == nullptr); - const bool has_d_multi_tensor = (d_base == nullptr); int64_t a_first = 0; int64_t a_last = 0; - if (!has_a_multi_tensor) { + if (!a_is_discrete) { a_first = A_meta.first_dims ? A_meta.first_dims[idx] : A_meta.uniform_first; a_last = A_meta.last_dims ? A_meta.last_dims[idx] : A_meta.uniform_last; } @@ -1358,24 +1356,25 @@ __global__ void setup_grouped_gemm_kernel( int64_t d_last = D_meta.last_dims ? D_meta.last_dims[idx] : D_meta.uniform_last; // Compute offsets (from explicit array, cumulative from per-tensor dims, or uniform) - int64_t a_offset = has_a_multi_tensor ? 0 : compute_grouped_tensor_offset(A_meta, idx); + int64_t a_offset = a_is_discrete ? 0 : compute_grouped_tensor_offset(A_meta, idx); int64_t b_offset = compute_grouped_tensor_offset(B_meta, idx); int64_t c_offset = compute_grouped_tensor_offset(C_meta, idx); int64_t d_offset = compute_grouped_tensor_offset(D_meta, idx); // Compute data pointers - A_ptrs[idx] = has_a_multi_tensor ? a_multi_tensor_args.data_ptrs[idx] - : (a_base + (a_offset * a_bits_per_elem) / 8); - B_ptrs[idx] = b_base + (b_offset * b_bits_per_elem) / 8; - C_ptrs[idx] = - has_c_multi_tensor ? c_multi_tensor_args.data_ptrs[idx] : (c_base + c_offset * c_elem_size); - D_ptrs[idx] = - has_d_multi_tensor ? d_multi_tensor_args.data_ptrs[idx] : (d_base + d_offset * d_elem_size); + A_ptrs[idx] = a_is_discrete + ? a_multi_tensor_args.data_ptrs[idx] + : (a_base == nullptr ? nullptr : a_base + (a_offset * a_bits_per_elem) / 8); + B_ptrs[idx] = b_base == nullptr ? nullptr : b_base + (b_offset * b_bits_per_elem) / 8; + C_ptrs[idx] = c_is_discrete ? c_multi_tensor_args.data_ptrs[idx] + : (c_base == nullptr ? nullptr : c_base + c_offset * c_elem_size); + D_ptrs[idx] = d_is_discrete ? d_multi_tensor_args.data_ptrs[idx] + : (d_base == nullptr ? nullptr : d_base + d_offset * d_elem_size); // Compute storage dimensions for cuBLAS matrix layouts from logical dims. // Rowwise and MXFP8 columnwise storage use logical row-major layout, viewed as // column-major rows=last, cols=first. Transposed columnwise storage reverses this. - if (has_a_multi_tensor) { + if (a_is_discrete) { a_rows[idx] = a_multi_tensor_args.rows[idx]; a_cols[idx] = a_multi_tensor_args.cols[idx]; } else if (a_storage_transposed) { @@ -1392,7 +1391,7 @@ __global__ void setup_grouped_gemm_kernel( b_rows[idx] = static_cast(b_last); b_cols[idx] = static_cast(b_first); } - if (has_d_multi_tensor) { + if (d_is_discrete) { d_rows[idx] = d_multi_tensor_args.rows[idx]; d_cols[idx] = d_multi_tensor_args.cols[idx]; } else { @@ -1407,7 +1406,7 @@ __global__ void setup_grouped_gemm_kernel( if (use_per_group_alpha_beta) { float a_amax_val = 0.0f; bool has_a_amax = false; - if (has_a_multi_tensor) { + if (a_is_discrete) { auto *a_amax_p = static_cast(a_multi_tensor_args.amax_ptrs[idx]); if (a_amax_p != nullptr) { a_amax_val = *a_amax_p; @@ -1475,10 +1474,12 @@ __global__ void setup_grouped_gemm_kernel( } }; - if (a_scale_base) { + if (a_is_discrete) { + a_scale_inv_ptrs[idx] = a_multi_tensor_args.scale_inv_ptrs[idx]; + } else if (a_scale_base) { fill_scale_ptr(a_scale_inv_ptrs, a_scale_base, A_meta, a_rowwise, a_scaling_mode); } else { - a_scale_inv_ptrs[idx] = a_multi_tensor_args.scale_inv_ptrs[idx]; + a_scale_inv_ptrs[idx] = nullptr; } if (b_scale_base) { fill_scale_ptr(b_scale_inv_ptrs, b_scale_base, B_meta, b_rowwise, b_scaling_mode); @@ -1495,7 +1496,7 @@ inline void launch_grouped_gemm_setup( const transformer_engine::Tensor *beta_tensor, bool use_per_group_alpha_beta, size_t num_tensors, cudaStream_t stream, const MultiTensorGroupGemmInputArgs &a_multi_tensor_args, const NVTETensor *C_list, - const NVTETensor *D_list, char *a_base, transformer_engine::DType c_dtype, + const NVTETensor *D_list, bool a_is_discrete, char *a_base, transformer_engine::DType c_dtype, transformer_engine::DType d_dtype) { // Use logical shape info from selection; storage transposes are tracked separately. TensorShapeInfo A_meta = A_sel.logical_tensor_shape; @@ -1503,31 +1504,31 @@ inline void launch_grouped_gemm_setup( TensorShapeInfo C_meta{}; TensorShapeInfo D_meta{}; - const bool has_d_multi_tensor = (D_list != nullptr); - const bool has_c_multi_tensor = (C_list != nullptr) || has_d_multi_tensor; + const bool d_is_discrete = (D_list != nullptr); + const bool c_is_discrete = (C_list != nullptr) || d_is_discrete; MultiTensorGroupGemmOutputArgs c_multi_tensor_args{}; MultiTensorGroupGemmOutputArgs d_multi_tensor_args{}; - if (has_d_multi_tensor) { + if (d_is_discrete) { d_multi_tensor_args = build_grouped_gemm_multi_out_args(D_list, num_tensors, num_tensors, d_dtype, "D"); } if (C_list != nullptr) { c_multi_tensor_args = build_grouped_gemm_multi_out_args(C_list, num_tensors, num_tensors, d_dtype, "C"); - } else if (has_d_multi_tensor) { + } else if (d_is_discrete) { c_multi_tensor_args = d_multi_tensor_args; } char *c_base = nullptr; char *d_base = nullptr; - if (!has_c_multi_tensor) { + if (!c_is_discrete) { NVTE_CHECK(C != nullptr && D != nullptr, "Grouped GEMM: C/D grouped tensors are required when no C list is provided"); C_meta = TensorShapeInfo::create_shape_info_for_C(C, D); c_base = static_cast(C->data.dptr); } - if (!has_d_multi_tensor) { + if (!d_is_discrete) { NVTE_CHECK(D != nullptr, "Grouped GEMM: D grouped tensor is required when no D list is provided"); D_meta = TensorShapeInfo::from_tensor(D); @@ -1548,8 +1549,8 @@ inline void launch_grouped_gemm_setup( const bool b_rowwise = B_sel.rowwise; // NVFP4 alpha needs A's amax from either A_sel.amax (grouped) or amax_ptrs (discrete). - const bool a_has_amax = (A_sel.amax != nullptr) || - (A_sel.dptr == nullptr && a_multi_tensor_args.amax_ptrs[0] != nullptr); + const bool a_has_amax = + a_is_discrete ? a_multi_tensor_args.amax_ptrs[0] != nullptr : A_sel.amax != nullptr; const bool needs_nvfp4_alpha = transformer_engine::is_nvfp_scaling(A_sel.scaling_mode) && a_has_amax && (B_sel.amax != nullptr); @@ -1558,11 +1559,12 @@ inline void launch_grouped_gemm_setup( ws.d_rows, ws.d_cols, ws.alpha_ptrs, ws.beta_ptrs, ws.a_scale_inv_ptrs, ws.b_scale_inv_ptrs, A_sel.dptr, B_sel.dptr, c_base, d_base, A_meta, B_meta, C_meta, D_meta, a_bits_per_elem, b_bits_per_elem, c_elem_size, d_elem_size, static_cast(alpha_tensor->data.dptr), - static_cast(beta_tensor->data.dptr), use_per_group_alpha_beta, - reinterpret_cast(A_sel.scale_inv), reinterpret_cast(B_sel.scale_inv), - a_rowwise, b_rowwise, A_sel.storage_transposed, B_sel.storage_transposed, A_sel.scaling_mode, - B_sel.scaling_mode, num_tensors, a_multi_tensor_args, c_multi_tensor_args, - d_multi_tensor_args, A_sel.amax ? static_cast(A_sel.amax) : nullptr, + static_cast(beta_tensor->data.dptr), use_per_group_alpha_beta, a_is_discrete, + c_is_discrete, d_is_discrete, reinterpret_cast(A_sel.scale_inv), + reinterpret_cast(B_sel.scale_inv), a_rowwise, b_rowwise, A_sel.storage_transposed, + B_sel.storage_transposed, A_sel.scaling_mode, B_sel.scaling_mode, num_tensors, + a_multi_tensor_args, c_multi_tensor_args, d_multi_tensor_args, + A_sel.amax ? static_cast(A_sel.amax) : nullptr, B_sel.amax ? static_cast(B_sel.amax) : nullptr, needs_nvfp4_alpha ? ws.nvfp4_computed_alpha : nullptr); @@ -1637,8 +1639,8 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT MultiTensorGroupGemmInputArgs a_multi_tensor_args{}; launch_grouped_gemm_setup(workspace.setup_workspace, A_sel, B_sel, inputC, outputD, alpha_tensor, beta_tensor, use_per_group_alpha_beta, num_tensors, stream, - a_multi_tensor_args, /*C_list=*/nullptr, /*D_list=*/nullptr, A_sel.dptr, - inputC->dtype(), outputD->dtype()); + a_multi_tensor_args, /*C_list=*/nullptr, /*D_list=*/nullptr, + /*a_is_discrete=*/false, A_sel.dptr, inputC->dtype(), outputD->dtype()); // Compute average dimensions for heuristics // K dimension: if transa, K is A's last dim; if not, K is A's first dim @@ -1792,8 +1794,8 @@ void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num launch_grouped_gemm_setup(workspace.setup_workspace, A_sel, B_sel, inputC, outputD, alpha_tensor, beta_tensor, use_per_group_alpha_beta, num_tensors, stream, - a_multi_tensor_args, /*C_list=*/nullptr, /*D_list=*/nullptr, nullptr, - inputC->dtype(), outputD->dtype()); + a_multi_tensor_args, /*C_list=*/nullptr, /*D_list=*/nullptr, + /*a_is_discrete=*/true, nullptr, inputC->dtype(), outputD->dtype()); GroupedGemmConfig gemm_config; gemm_config.use_split_accumulator = config_.use_split_accumulator; @@ -1877,8 +1879,8 @@ void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, MultiTensorGroupGemmInputArgs a_multi_tensor_args{}; launch_grouped_gemm_setup(workspace.setup_workspace, A_sel, B_sel, /*C=*/nullptr, /*D=*/nullptr, alpha_tensor, beta_tensor, use_per_group_alpha_beta, num_tensors, - stream, a_multi_tensor_args, C_list, D_list, A_sel.dptr, d_dtype, - d_dtype); + stream, a_multi_tensor_args, C_list, D_list, + /*a_is_discrete=*/false, A_sel.dptr, d_dtype, d_dtype); GroupedGemmConfig gemm_config; gemm_config.use_split_accumulator = config_.use_split_accumulator; @@ -1913,8 +1915,6 @@ void launch_grouped_bias_add(const transformer_engine::GroupedTensor *outputD, NVTE_CHECK(outputD->num_tensors >= 1, api_name, ": number of tensors must be at least 1"); NVTE_CHECK(outputD->num_tensors == bias_tensor->num_tensors, api_name, ": output and bias must have the same number of tensors"); - NVTE_CHECK(outputD->has_data(), api_name, ": output is missing row-wise data"); - NVTE_CHECK(bias_tensor->has_data(), api_name, ": bias is missing row-wise data"); NVTE_CHECK(outputD->dtype() == bias_tensor->dtype(), api_name, ": output and bias must have matching dtypes"); NVTE_CHECK(bias_tensor->all_same_first_dim(), api_name, @@ -1925,15 +1925,23 @@ void launch_grouped_bias_add(const transformer_engine::GroupedTensor *outputD, NVTE_CHECK(outputD->get_common_last_dim() == bias_tensor->get_common_last_dim(), api_name, ": output and bias last dims must match"); + const int num_tensors = static_cast(outputD->num_tensors); + NVTE_CHECK(num_tensors <= kMaxGroups, api_name, " supports at most ", kMaxGroups, + " tensors, got ", num_tensors); + const int total_rows = static_cast(outputD->logical_shape.data[0]); + // A valid zero-sized CUDA allocation may have a null data pointer. + if (total_rows == 0) { + return; + } + + NVTE_CHECK(outputD->has_data(), api_name, ": output is missing row-wise data"); + NVTE_CHECK(bias_tensor->has_data(), api_name, ": bias is missing row-wise data"); + const TensorShapeInfo d_meta = TensorShapeInfo::from_tensor(outputD); const DType dtype = outputD->dtype(); constexpr int kThreads = 128; - const int num_tensors = static_cast(outputD->num_tensors); - NVTE_CHECK(num_tensors <= kMaxGroups, api_name, " supports at most ", kMaxGroups, - " tensors, got ", num_tensors); - const int total_rows = static_cast(outputD->logical_shape.data[0]); const int n = static_cast(outputD->get_common_last_dim()); const size_t elem_size = typeToSize(dtype); @@ -2001,8 +2009,6 @@ void nvte_grouped_scaled_bias_add(const NVTEGroupedTensor output, const NVTEGrou const GroupedTensor *bias_tensor = convertNVTEGroupedTensorCheck(bias); const Tensor *scale_tensor = convertNVTETensorCheck(scale); - NVTE_CHECK(scale_tensor->data.dptr != nullptr, - "Grouped scaled bias add: scale tensor must not be null"); NVTE_CHECK(scale_tensor->dtype() == DType::kFloat32, "Grouped scaled bias add: scale must be float32"); NVTE_CHECK(scale_tensor->data.shape.size() == 1, @@ -2011,6 +2017,10 @@ void nvte_grouped_scaled_bias_add(const NVTEGroupedTensor output, const NVTEGrou const size_t total_rows = static_cast(outputD->logical_shape.data[0]); NVTE_CHECK(scale_tensor->data.shape[0] == total_rows, "Grouped scaled bias add: scale size (", scale_tensor->data.shape[0], ") must equal total rows (", total_rows, ")"); + if (total_rows > 0) { + NVTE_CHECK(scale_tensor->data.dptr != nullptr, + "Grouped scaled bias add: scale tensor must not be null"); + } const float *scale_ptr = static_cast(scale_tensor->data.dptr); launch_grouped_bias_add(outputD, bias_tensor, scale_ptr, true, stream); diff --git a/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu index d5dbd0bd82..0f2456c975 100644 --- a/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu @@ -13,8 +13,6 @@ #include #include -#include -#include #include #include "common/common.h" @@ -746,7 +744,7 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device_g cutlass::arch::NamedBarrier::sync(NumEpilogueColQuantThreadCount, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); - // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} + // Aligning with TensorEngine's recipe to generate scale factors static constexpr float fp4_max = 6.0f; static constexpr float fp8_max = 448.0f; static constexpr float fp4_max_inv = 1.0f / fp4_max; @@ -1003,7 +1001,7 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device_g shape_rep, num_tensors, (scheduler.tile_n_base() * size<1>(epilogue_tiler)) * M, packed_N, M, offsets); float a_global_amax_val = shared_storage.global_a_amax[group_idx]; - // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} + // Aligning with TensorEngine's recipe to generate scale factors static constexpr float fp4_max = 6.0f; static constexpr float fp8_max = 448.0f; static constexpr float fp4_max_inv = 1.0f / fp4_max; diff --git a/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu index e2325dd0fc..4b1435f9eb 100644 --- a/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu @@ -13,8 +13,6 @@ #include #include -#include -#include #include #include "common/common.h" diff --git a/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu index cab0b38589..2e6d383ce1 100644 --- a/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu @@ -13,8 +13,6 @@ #include #include -#include -#include #include #include "common/common.h" @@ -728,7 +726,7 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device( cutlass::arch::NamedBarrier::sync(NumEpilogueColQuantThreadCount, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); - // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} + // Aligning with TensorEngine's recipe to generate scale factors static constexpr float fp4_max = 6.0f; static constexpr float fp8_max = 448.0f; static constexpr float fp4_max_inv = 1.0f / fp4_max; @@ -980,7 +978,7 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device( int group_idx = GetGroupIdx(&args, scheduler.tile_n_base() * size<1>(epilogue_tiler)); float a_global_amax_val = shared_storage.global_a_amax[group_idx]; - // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} + // Aligning with TensorEngine's recipe to generate scale factors static constexpr float fp4_max = 6.0f; static constexpr float fp8_max = 448.0f; static constexpr float fp4_max_inv = 1.0f / fp4_max; diff --git a/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu index 50b9f63bdd..433da1f0f0 100644 --- a/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu @@ -13,8 +13,6 @@ #include #include -#include -#include #include #include "common/common.h" diff --git a/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu index 479922a9bf..8d8ab20165 100644 --- a/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu @@ -13,8 +13,6 @@ #include #include -#include -#include #include #include "common/common.h" @@ -709,7 +707,7 @@ __global__ static void row_col_rht_gemm_device( auto thr_t2r = tiled_t2r.get_slice(local_thread_idx); auto thr_r2g = tiled_r2g.get_slice(local_thread_idx); - // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} + // Aligning with TensorEngine's recipe to generate scale factors static constexpr float fp4_max = 6.0f; static constexpr float fp8_max = 448.0f; float const fp4_max_inv = 1.0f / fp4_max; @@ -905,7 +903,7 @@ __global__ static void row_col_rht_gemm_device( cute::Tensor tQArSFA = make_tensor_like(tQAgSFA(_, _, _, _0{}, _0{})); cute::Tensor tQApSFA = thr_s2r.partition_D(pSFA_mn); - // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} + // Aligning with TensorEngine's recipe to generate scale factors static constexpr float fp4_max = 6.0f; static constexpr float fp8_max = 448.0f; float const fp4_max_inv = 1.0f / fp4_max; diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index 554d8c1ac9..4d6d24ba65 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -161,10 +161,12 @@ void nvte_quantize_dbias(const NVTETensor input, NVTETensor output, NVTETensor d * \param[in,out] output Output grouped FP8/MXFP8 tensor. * \param[out] dbias Result of the reduction of the input along columns. * \param[out] workspace Workspace tensor. + * \param[in] quant_config Quantization configuration. * \param[in] stream CUDA stream used for the operation. */ void nvte_group_quantize_dbias(const NVTEGroupedTensor input, NVTEGroupedTensor output, - NVTEGroupedTensor dbias, NVTETensor workspace, cudaStream_t stream); + NVTEGroupedTensor dbias, NVTETensor workspace, + const NVTEQuantizationConfig quant_config, cudaStream_t stream); /*! \brief Computes backward of GeLU operation on the input, then casts to FP8/MXFP8. * Additionally, reduces the result of the GeLU backward along columns. diff --git a/transformer_engine/common/include/transformer_engine/comm_window.h b/transformer_engine/common/include/transformer_engine/comm_window.h index 424c350bbd..ed67774493 100644 --- a/transformer_engine/common/include/transformer_engine/comm_window.h +++ b/transformer_engine/common/include/transformer_engine/comm_window.h @@ -26,6 +26,9 @@ struct ncclWindow_vidmem; typedef struct { struct ncclWindow_vidmem* window; /*!< NCCL window, or NULL to use the raw data pointer. */ uint64_t offset; /*!< Byte offset of the payload within window. */ + struct ncclWindow_vidmem* + scale_window; /*!< Window for a block-scaled tensor's scale-inverse, or NULL for raw. */ + uint64_t scale_offset; /*!< Byte offset of the scale-inverse within scale_window. */ } NVTECommWindow; #ifdef __cplusplus diff --git a/transformer_engine/common/include/transformer_engine/ep.h b/transformer_engine/common/include/transformer_engine/ep.h index 224622fd41..31417fbff1 100644 --- a/transformer_engine/common/include/transformer_engine/ep.h +++ b/transformer_engine/common/include/transformer_engine/ep.h @@ -44,7 +44,8 @@ typedef struct { int num_experts; /*! Upper bound on tokens this rank sends per dispatch. */ int max_tokens_per_rank; - /*! Upper bound on tokens this rank receives per dispatch (must be > 0). */ + /*! Upper bound on tokens this rank receives per dispatch. 0 selects eager + * mode: the caller sizes recv buffers to the per-routing recv count. */ int max_recv_tokens_per_rank; /*! Token hidden dimension. */ int hidden_dim; @@ -58,6 +59,12 @@ typedef struct { * by NVTECommWindow handles and transfer in place (no staging copies); * 0 (default) = staged. */ int zero_copy; + /*! Per-token top-k; sizes NCCL EP internal buffers. Required in eager mode + * (max_recv_tokens_per_rank == 0); 0 = unset. */ + int num_topk; + /*! Recv overflow policy. Nonzero drops tokens past max_recv_tokens_per_rank + * and continues; 0 (default) traps. Not supported in eager mode. */ + int drop_on_overflow; } NVTEEpGroupConfig; /*! \brief Per-layer configuration consumed by nvte_ep_handle_mem_size and @@ -121,13 +128,14 @@ size_t nvte_ep_handle_mem_size(const NVTEEpLayerConfig* layer_cfg); * AllGathers topk_idx across the EP group and stages per-expert offsets and * counts into handle_mem so the matching dispatch/combine/_bwd can run with * no further routing computation. Must precede every dispatch/combine/_bwd - * that uses this handle_mem. recv_tokens_per_expert becomes host-valid after a - * stream sync. + * that uses this handle_mem. recv_tokens_per_expert and total_recv_tokens_per_rank + * become host-valid after a stream sync. * * \param[in] handle_mem uint8 routing-state buffer. * \param[in] topk_idx [T, top_k] int64 routing indices. - * \param[out] recv_tokens_per_expert [num_local_experts] int32 counts. - * \param[out] total_recv_tokens_per_rank Reserved placeholder; may be null. Unused for now. + * \param[out] recv_tokens_per_expert [num_local_experts] int32/int64 counts. + * \param[out] total_recv_tokens_per_rank Optional [1] int32/int64 scalar: padded + * recv-slot total for this rank. May be null. * \param[in] layer_cfg Per-call layer configuration (struct_size set). * \param[in] stream CUDA stream. */ @@ -142,6 +150,12 @@ void nvte_ep_prepare(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor recv * *_win arguments enable zero-copy via symmem windows; pass NVTECommWindow{} * when unused. Requires a prior nvte_ep_prepare on this handle_mem. * + * tokens/recv_tokens may be high-precision (bf16/fp16) or FP8: + * for the latter, set rowwise data and rowwise scale-inverse (unswizzled + * [T, hidden_dim/block]) on the tensor and the scales are routed alongside the + * data. tokens and recv_tokens must share a scaling mode. For now, only MXFP8 + * (NVTE_MXFP8_1D_SCALING, e4m3 data + e8m0 scales) is supported. + * * \param[in] handle_mem uint8 routing-state buffer (from prepare). * \param[in] topk_idx [T, top_k] int64 sparse routing indices. * \param[in] tokens [T, hidden_dim] input tokens. diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 39dbe12165..ac8b1e5c0f 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -536,16 +536,22 @@ void nvte_cp_thd_out_correction(NVTETensor out, const NVTETensor &out_per_step, const NVTETensor &cu_seqlens, int only_second_half, int lse_packed, cudaStream_t stream); -/*! \brief Correct the THD format output of context parallelism in forward pass. +/*! \brief Update the two halves of each packed THD sequence during context-parallel backward. * * \warning This API is **experimental** and subject to change. * - * \param[out] grad Output tensor. - * \param[in] grad_per_step THD format gradient of context parallelism. - * \param[in] cu_seqlens Cumulative sequence lengths, [batch_size + 1]. - * \param[in] first_half One of ("add", "copy", "none") correction op for first half. - * \param[in] second_half One of ("add", "copy", "none") correction op for second half. - Must be different from first_half. + * first_half and second_half control how grad is updated from grad_per_step: "add" accumulates, + * "copy" replaces, "none" preserves, and "zero" clears. FP16, BF16, and FP32 gradients support + * (add, none), (none, add), (copy, none), (none, copy), (copy, zero), (zero, copy), (add, copy), + * and (copy, add). FP8 gradients are stored as raw encoded bytes, which this kernel cannot add + * numerically. They support only (copy, zero) and (zero, copy): copy preserves the FP8 values and + * zero clears the inactive sequence half. + * + * \param[in,out] grad Packed THD gradient to update. + * \param[in] grad_per_step Gradient from the current context-parallel step. + * \param[in] cu_seqlens Packed-sequence boundaries, [batch_size + 1]. + * \param[in] first_half Operation for each sequence's first half. + * \param[in] second_half Operation for each sequence's second half. * \param[in] stream CUDA stream used for this operation. */ void nvte_cp_thd_grad_correction(NVTETensor grad, const NVTETensor &grad_per_step, diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index bb87e7be35..6f3a79277e 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -447,9 +447,11 @@ enum NVTEQuantizationConfigAttribute { * of ordinary NVFP4 fast-math settings. */ kNVTEQuantizationConfigNVFP44Over6ErrUseFastMath = 9, + /*! Whether to use 2D block scaling for MXFP8 */ + kNVTEQuantizationConfigMXFP82DQuantization = 10, #ifdef USE_ROCM /*! Whether to apply Hadamard transform before MXFP4 quantization */ - kNVTEQuantizationConfigMXFP4UseHadamard = 10, + kNVTEQuantizationConfigMXFP4UseHadamard = 11, #endif kNVTEQuantizationConfigNumAttributes }; @@ -1558,6 +1560,13 @@ class QuantizationConfigWrapper { &val, sizeof(val)); } + /*! \brief Set whether to use 2D block scaling for MXFP8 */ + void set_mxfp8_2d_quantization(bool mxfp8_2d_quantization) { + const auto val = static_cast(mxfp8_2d_quantization); + nvte_set_quantization_config_attribute(config_, kNVTEQuantizationConfigMXFP82DQuantization, + &val, sizeof(val)); + } + /*! \brief Set whether to use stochastic rounding */ void set_stochastic_rounding(bool stochastic_rounding) { const auto val = static_cast(stochastic_rounding); diff --git a/transformer_engine/common/newton_schulz/newton_schulz.cpp b/transformer_engine/common/newton_schulz/newton_schulz.cpp index 0d6426a156..5eeaf2da00 100644 --- a/transformer_engine/common/newton_schulz/newton_schulz.cpp +++ b/transformer_engine/common/newton_schulz/newton_schulz.cpp @@ -134,6 +134,19 @@ void FreeWorkspace(NVTECusolverMpCtx* ctx) { NVTECusolverMpCtx* nvte_cusolvermp_ctx_create(ncclComm_t comm, int nranks, int rank) { NVTE_API_CALL(nvte_cusolvermp_ctx_create); + NVTE_CHECK(comm != nullptr, "NCCL communicator must be non-null"); + NVTE_CHECK(nranks > 0, "Number of ranks must be positive, got ", nranks); + NVTE_CHECK(rank >= 0 && rank < nranks, "Rank ", rank, " is outside [0, ", nranks, ")"); + + int comm_nranks{}; + int comm_rank{}; + NVTE_CHECK_NCCL(ncclCommCount(comm, &comm_nranks)); + NVTE_CHECK_NCCL(ncclCommUserRank(comm, &comm_rank)); + NVTE_CHECK(comm_nranks == nranks, "NCCL communicator has ", comm_nranks, + " ranks, but the process group reports ", nranks); + NVTE_CHECK(comm_rank == rank, "NCCL communicator rank is ", comm_rank, + ", but the process group reports ", rank); + int device_id{}; NVTE_CHECK_CUDA(cudaGetDevice(&device_id)); diff --git a/transformer_engine/common/normalization/common.cpp b/transformer_engine/common/normalization/common.cpp index e6de388ce8..0325c84678 100644 --- a/transformer_engine/common/normalization/common.cpp +++ b/transformer_engine/common/normalization/common.cpp @@ -361,7 +361,17 @@ CudnnNormalizationPlan::CudnnNormalizationPlan(NVTE_Norm_Type NormType, NVTE_Nor if (_training) _rsigma->set_output(true).set_data_type(get_cudnn_fe_dtype(ctype)); - const auto ZDtype = _fp8_out ? ctype : otype; + auto ZDtype = _fp8_out ? ctype : otype; + if (_fp8_out) { + const bool use_input_dtype = cudnnGetVersion() >= 92500 && _ndim_scale_block == 1 && + use_cudnn_mxfp8_norm_output_in_input_dtype(); + if (use_input_dtype) { + NVTE_WARN( + "The cuDNN MXFP8 normalization intermediate output uses the input dtype (itype) " + "instead of the compute dtype; otype still applies to the final quantized output."); + ZDtype = itype; + } + } _z->set_output(!_fp8_out).set_data_type(get_cudnn_fe_dtype(ZDtype)); if (_fp8_out) { @@ -633,6 +643,12 @@ bool& _zero_centered_gamma_in_weight_dtype() { bool& use_zero_centered_gamma_in_weight_dtype() { return _zero_centered_gamma_in_weight_dtype(); } #endif +bool use_cudnn_mxfp8_norm_output_in_input_dtype() { + static bool flag = + transformer_engine::getenv("NVTE_CUDNN_MXFP8_NORM_OUTPUT_IN_INPUT_DTYPE"); + return flag; +} + } // namespace normalization } // namespace transformer_engine diff --git a/transformer_engine/common/normalization/common.h b/transformer_engine/common/normalization/common.h index 07670ee65f..d574bc4f8a 100644 --- a/transformer_engine/common/normalization/common.h +++ b/transformer_engine/common/normalization/common.h @@ -421,6 +421,7 @@ bool use_cudnn_norm_fwd(); bool use_cudnn_norm_bwd(); bool& use_zero_centered_gamma_in_weight_dtype(); +bool use_cudnn_mxfp8_norm_output_in_input_dtype(); #endif #ifdef __HIP_PLATFORM_AMD__ @@ -464,7 +465,7 @@ void rocm_norm_mxfp8_quantize(LaunchParams &launch_params) ); ); } -#endif +#endif } // namespace normalization } // namespace transformer_engine diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 35081cbf4b..b32cb6c720 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -389,6 +389,8 @@ class MXFP8BlockScaling(Recipe): `high_precision` keeps original high-precision operands for backward, and `dequantized` dequantizes saved operands to the active high-precision compute dtype (e.g. BF16/FP16/FP32) for backward. + enable_2d_quantization : bool, default = False + If set to `True`, 2D block scaling is used for weight tensors. """ margin: int = 0 @@ -396,6 +398,7 @@ class MXFP8BlockScaling(Recipe): fp8_dpa: bool = False fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) + enable_2d_quantization: bool = False def __post_init__(self) -> None: assert self.fp8_format != Format.E5M2, "Pure E5M2 training is not supported." @@ -408,7 +411,8 @@ def _make_repr(self) -> str: f"recipe_type={self.__class__.__name__}, " f"margin={self.margin}, " f"format={str(self.fp8_format).split('.')[1]}, " - f"backward_override={self.backward_override}" + f"backward_override={self.backward_override}, " + f"enable_2d_quantization={self.enable_2d_quantization}" ) @@ -432,6 +436,12 @@ class Float8BlockScaling(Recipe): NOTE: To relax the default constraint that scales be powers of 2, set env variable NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1 to override it for the recipe defaults. + NOTE: FP8 block scaling requires split accumulation for numerical accuracy, so + ``fp8_gemm_fprop``/``fp8_gemm_dgrad``/``fp8_gemm_wgrad`` all fix + ``use_split_accumulator=True`` (enforced in ``__post_init__``). The fused grouped + GEMM path (GroupedLinear) always uses split accumulation for FP8 block scaling and + ignores any caller- or recipe-supplied ``use_split_accumulator`` value. + Parameters ---------- fp8_format : {Format.E4M3, Format.HYBRID}, default = Format.E4M3 @@ -688,6 +698,14 @@ class CustomRecipe(Recipe): `high_precision` keeps original high-precision operands for backward, and `dequantized` dequantizes saved operands to the active high-precision compute dtype (e.g. BF16/FP16/FP32) for backward. + quantization_alignment : int, default = 128 + Conservative recipe-wide fallback used by automatic padding for grouped operations. + This must be at least the largest alignment required by any quantizer + that ``qfactory`` may return for a grouped operation, across all roles + and module names. The default of 128 safely supports all current TE + formats. It can be lowered when the factory's full output space is known; + for example, a factory restricted to MXFP8 may use 32. + Automatic padding reads this value without invoking ``qfactory``. """ qfactory: Callable[..., Any] @@ -700,17 +718,21 @@ class CustomRecipe(Recipe): fp8_dpa: bool = False fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) + quantization_alignment: int = 128 def __post_init__(self) -> None: assert ( self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." + if self.quantization_alignment <= 0: + raise ValueError("CustomRecipe quantization_alignment must be positive.") def _make_repr(self) -> str: return ( f"recipe_type={self.__class__.__name__}, " f"qfactory={self.qfactory}, " - f"backward_override={self.backward_override}" + f"backward_override={self.backward_override}, " + f"quantization_alignment={self.quantization_alignment}" ) @dataclass() diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index c2016cc5da..469700eede 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -11,9 +11,13 @@ #include #include #include +#include #include +#include +#include #include #include +#include #include #include #include @@ -449,6 +453,47 @@ void CheckOutputGroupedTensor(const GroupedTensor &t, std::string_view name, boo CheckGroupedTensorShapeArrays(t, name); } +namespace { + +constexpr size_t kDefaultTensorHandlePoolSizeMB = 20; +constexpr size_t kBytesPerMB = 1024 * 1024; + +size_t GetTensorHandlePoolSizeMB(const char *env_var) { + const char *env_value = std::getenv(env_var); + if (env_value == nullptr || env_value[0] == '\0') { + return kDefaultTensorHandlePoolSizeMB; + } + + const std::string value(env_value); + constexpr const char *kWhitespace = " \t\n\r\f\v"; + const size_t first = value.find_first_not_of(kWhitespace); + const size_t last = value.find_last_not_of(kWhitespace); + NVTE_CHECK(first != std::string::npos, env_var, " must be a positive integer."); + + const char *begin = value.c_str() + first; + const char *expected_end = value.c_str() + last + 1; + errno = 0; + char *end = nullptr; + const uint64_t pool_size_mb = std::strtoul(begin, &end, 10); + + NVTE_CHECK(value[first] >= '0' && value[first] <= '9' && end == expected_end && pool_size_mb > 0, + env_var, " must be a positive integer, got \"", value, "\"."); + NVTE_CHECK(errno != ERANGE, env_var, " is too large."); + NVTE_CHECK(pool_size_mb <= std::numeric_limits::max() / kBytesPerMB, env_var, + " is too large."); + return static_cast(pool_size_mb); +} + +size_t GetTensorHandlePoolCapacity(size_t pool_size_mb, size_t handle_size, const char *handle_name, + const char *env_var) { + const size_t pool_size_bytes = pool_size_mb * kBytesPerMB; + NVTE_CHECK(pool_size_bytes >= handle_size, env_var, "=", pool_size_mb, + " MiB is too small for one ", handle_name, " handle of size ", handle_size, " bytes."); + return pool_size_bytes / handle_size; +} + +} // namespace + class TensorAllocator { public: static TensorAllocator &instance() { @@ -462,8 +507,10 @@ class TensorAllocator { std::lock_guard lock(mutex); const size_t available = free_list.size() + (memory.capacity() - memory.size()); NVTE_CHECK(available >= N, "Cannot allocate ", N, - " new NVTETensors. Maximum number of tensors reached: ", MAX_TENSOR_NUM, - ". There is probably a memory leak in your application."); + " new NVTETensors. Maximum number of tensors reached: ", MAX_TENSOR_NUM, " (", + TENSOR_HANDLE_POOL_SIZE_MB, + " MiB handle pool). If your application legitimately needs more tensor handles, " + "increase NVTE_TENSOR_HANDLE_POOL_SIZE_MB."); for (size_t i = 0; i < N; ++i) { uintptr_t index; if (!free_list.empty()) { @@ -535,9 +582,11 @@ class TensorAllocator { std::mutex mutex; std::atomic size; - // Allocate at most 20 MB for tensors // Should be replaced by virtual memory allocation - const size_t MAX_TENSOR_NUM = 20 * 1024 * 1024 / sizeof(Tensor); + const size_t TENSOR_HANDLE_POOL_SIZE_MB = + GetTensorHandlePoolSizeMB("NVTE_TENSOR_HANDLE_POOL_SIZE_MB"); + const size_t MAX_TENSOR_NUM = GetTensorHandlePoolCapacity( + TENSOR_HANDLE_POOL_SIZE_MB, sizeof(Tensor), "NVTETensor", "NVTE_TENSOR_HANDLE_POOL_SIZE_MB"); std::vector free_list; std::vector memory; bool debug = false; @@ -588,7 +637,9 @@ class GroupedTensorAllocator { } NVTE_ERROR( "Cannot allocate a new NVTEGroupedTensor. Maximum number of grouped tensors reached: ", - MAX_GROUPED_TENSOR_NUM, ". There is probably a memory leak in your application."); + MAX_GROUPED_TENSOR_NUM, " (", GROUPED_TENSOR_HANDLE_POOL_SIZE_MB, + " MiB handle pool). If your application legitimately needs more grouped tensor handles, " + "increase NVTE_GROUPED_TENSOR_HANDLE_POOL_SIZE_MB."); } void Free(NVTEGroupedTensor t) { @@ -620,8 +671,11 @@ class GroupedTensorAllocator { std::mutex mutex; std::atomic size; - // Allocate at most 20 MB for grouped tensors - const size_t MAX_GROUPED_TENSOR_NUM = 20 * 1024 * 1024 / sizeof(GroupedTensor); + const size_t GROUPED_TENSOR_HANDLE_POOL_SIZE_MB = + GetTensorHandlePoolSizeMB("NVTE_GROUPED_TENSOR_HANDLE_POOL_SIZE_MB"); + const size_t MAX_GROUPED_TENSOR_NUM = + GetTensorHandlePoolCapacity(GROUPED_TENSOR_HANDLE_POOL_SIZE_MB, sizeof(GroupedTensor), + "NVTEGroupedTensor", "NVTE_GROUPED_TENSOR_HANDLE_POOL_SIZE_MB"); std::vector free_list; std::vector memory; }; @@ -1122,6 +1176,9 @@ void nvte_get_quantization_config_attribute(NVTEQuantizationConfig config, case kNVTEQuantizationConfigNVFP44Over6ErrUseFastMath: bool_to_uint8(config_.nvfp4_4over6_err_use_fast_math, buf); break; + case kNVTEQuantizationConfigMXFP82DQuantization: + bool_to_uint8(config_.mxfp8_2d_quantization, buf); + break; #ifdef __HIP_PLATFORM_AMD__ case kNVTEQuantizationConfigMXFP4UseHadamard: bool_to_uint8(config_.mxfp4_use_hadamard, buf); @@ -1194,6 +1251,9 @@ void nvte_set_quantization_config_attribute(NVTEQuantizationConfig config, case kNVTEQuantizationConfigNVFP44Over6ErrUseFastMath: uint8_to_bool(buf, config_.nvfp4_4over6_err_use_fast_math); break; + case kNVTEQuantizationConfigMXFP82DQuantization: + uint8_to_bool(buf, config_.mxfp8_2d_quantization); + break; #ifdef __HIP_PLATFORM_AMD__ case kNVTEQuantizationConfigMXFP4UseHadamard: uint8_to_bool(buf, config_.mxfp4_use_hadamard); diff --git a/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu b/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu index 73d86f19d4..7826c87718 100644 --- a/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu +++ b/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu @@ -175,13 +175,15 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK) static_assert(std::is_same::value); const CType scale_inv = 1.0f / block_tile_scale; - size_t row_idx = tile_id_y; - size_t col_idx = tile_id_x; - tile_scales_inv_c[row_idx * scale_stride_y + col_idx * scale_stride_x] = scale_inv; + if (tile_scales_inv_c != nullptr) { + size_t row_idx = tile_id_y; + size_t col_idx = tile_id_x; + tile_scales_inv_c[row_idx * scale_stride_y + col_idx * scale_stride_x] = scale_inv; + } if constexpr (kReturnTranspose) { - row_idx = tile_id_x; - col_idx = tile_id_y; + size_t row_idx = tile_id_x; + size_t col_idx = tile_id_y; tile_scales_inv_t[row_idx * scale_t_stride_y + col_idx * scale_t_stride_x] = scale_inv; } } @@ -203,7 +205,9 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK) thrd_tile_out_trans[j].data.elt[i] = scaled_elt; } } - tmp_output_c.store_to(output_c + thread_tile_start_idx + i * row_length); + if (output_c != nullptr) { + tmp_output_c.store_to(output_c + thread_tile_start_idx + i * row_length); + } } // Step 4: store transpose into shared memory @@ -407,13 +411,15 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK) block_scaled_cast_transpose static_assert(std::is_same::value); const CType scale_inv = 1.0f / block_tile_scale; - size_t row_idx = tile_id_y; - size_t col_idx = tile_id_x; - tile_scales_inv_c[row_idx * scale_stride_y + col_idx * scale_stride_x] = scale_inv; + if (tile_scales_inv_c != nullptr) { + size_t row_idx = tile_id_y; + size_t col_idx = tile_id_x; + tile_scales_inv_c[row_idx * scale_stride_y + col_idx * scale_stride_x] = scale_inv; + } if constexpr (kReturnTranspose) { - row_idx = tile_id_x; - col_idx = tile_id_y; + size_t row_idx = tile_id_x; + size_t col_idx = tile_id_y; tile_scales_inv_t[row_idx * scale_t_stride_y + col_idx * scale_t_stride_x] = scale_inv; } } @@ -452,8 +458,10 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK) block_scaled_cast_transpose thrd_tile_out_trans[j].data.elt[i] = scaled_elt; } } - tmp_output_c.store_to_elts(output_c + thread_tile_start_idx + i * row_length, 0, - thread_tile_ncols); + if (output_c != nullptr) { + tmp_output_c.store_to_elts(output_c + thread_tile_start_idx + i * row_length, 0, + thread_tile_ncols); + } } if constexpr (kReturnTranspose) { @@ -513,19 +521,26 @@ void quantize_transpose_square_blockwise(const SimpleTensor& input, SimpleTensor "with MXFP8, which requires using power of two scaling factors."); } - NVTE_CHECK(input.shape == output.shape, "Input and output must have the same shape."); + const bool return_identity = output.dptr != nullptr; + if (return_identity) { + NVTE_CHECK(input.shape == output.shape, "Input and output must have the same shape."); + } + NVTE_CHECK(return_identity || return_transpose, + "At least one of rowwise or columnwise output must be requested."); const size_t row_length = input.shape.size() > 0 ? input.shape.back() : 1; size_t num_rows = 1; for (size_t i = 0; (i < input.shape.size() - 1) && (input.shape.size() > 0); ++i) { num_rows *= input.shape.at(i); } - NVTE_CHECK(scale_inv.shape.size() == 2, "scale_inv must have 2 dimensions."); - - size_t scale_k = scale_inv.shape[1]; - - const size_t scale_stride_x = 1; - const size_t scale_stride_y = scale_k; + size_t scale_k = 0; + const size_t scale_stride_x = return_identity ? 1 : 0; + size_t scale_stride_y = 0; + if (return_identity) { + NVTE_CHECK(scale_inv.shape.size() == 2, "scale_inv must have 2 dimensions."); + scale_k = scale_inv.shape[1]; + scale_stride_y = scale_k; + } size_t scale_t_stride_x = 0; size_t scale_t_stride_y = 0; @@ -543,7 +558,9 @@ void quantize_transpose_square_blockwise(const SimpleTensor& input, SimpleTensor ") and output_t (shape=", output_t.shape, ") have incompatible dims."); } } - NVTE_CHECK(output.dtype == output_t.dtype, "output and output_t need to have the same type."); + if (return_identity) { + NVTE_CHECK(output.dtype == output_t.dtype, "output and output_t need to have the same type."); + } NVTE_CHECK(scale_inv_t.shape.size() == 2, "scale_inv_t must have 2 dimensions."); @@ -551,6 +568,8 @@ void quantize_transpose_square_blockwise(const SimpleTensor& input, SimpleTensor scale_t_stride_y = scale_inv_t.shape[1]; } + const auto out_dtype = return_identity ? output.dtype : output_t.dtype; + const size_t num_blocks_x = DIVUP(row_length, BLOCK_TILE_DIM); const size_t num_blocks_y = DIVUP(num_rows, BLOCK_TILE_DIM); @@ -558,7 +577,7 @@ void quantize_transpose_square_blockwise(const SimpleTensor& input, SimpleTensor input.dtype, InputType, TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( - output.dtype, OutputType, + out_dtype, OutputType, TRANSFORMER_ENGINE_SWITCH_CONDITION( return_transpose, kReturnTranspose, diff --git a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu index a67ad85812..e1e0433214 100644 --- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu +++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu @@ -416,11 +416,11 @@ __device__ __forceinline__ __nv_fp4x4_e2m1 cvt_fp32_to_fp4_4x_with_stochastic_ro : "f"(in01.y), "f"(in01.x), "f"(in23.y), "f"(in23.x), "r"(rbits)); return *reinterpret_cast<__nv_fp4x4_e2m1*>(&out_4x); } else { - NVTE_DEVICE_ERROR( - "FP4 cvt.rs PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); - uint16_t dummy = 0; - return *reinterpret_cast<__nv_fp4x4_e2m1*>(&dummy); + const float q0 = ptx::stochastic_round_fp4_e2m1(in01.x, rbits); + const float q1 = ptx::stochastic_round_fp4_e2m1(in01.y, rbits >> 8); + const float q2 = ptx::stochastic_round_fp4_e2m1(in23.x, rbits >> 16); + const float q3 = ptx::stochastic_round_fp4_e2m1(in23.y, rbits >> 24); + return __nv_fp4x4_e2m1(make_float4(q0, q1, q2, q3)); } #else // It is like ptx.cuh::mul_cvt_fp32_to_fp4_4x_with_stochastic_rounding but w/o scaling diff --git a/transformer_engine/common/triton/mhc.py b/transformer_engine/common/triton/mhc.py index 965bb437ff..ddad878e1a 100644 --- a/transformer_engine/common/triton/mhc.py +++ b/transformer_engine/common/triton/mhc.py @@ -12,6 +12,12 @@ import triton import triton.language as tl +MAX_GRID_DIM_Y = 65535 # Maximum grid dimension in Y direction for current CUDA architectures + + +def align_to(x, alignment): + return ((x + alignment - 1) // alignment) * alignment + def projection_config_fwd(): block_m = [64, 128] @@ -29,28 +35,50 @@ def projection_config_fwd(): num_stages=s, ) ) - if os.environ.get("NVTE_DISABLE_TRITON_AUTOTUNING", "0") == "1": - configs = configs[:1] return configs -def projection_config_bwd(): - block_m = [32, 128] - block_k = [128] - warps = [2] - stages = [2, 3, 4] - - configs = [] - for m, bk, w, s in itertools.product(block_m, block_k, warps, stages): - configs.append( - triton.Config({"BLOCK_SIZE_M": m, "BLOCK_SIZE_K": bk}, num_warps=w, num_stages=s) - ) +def projection_prune_fwd(configs, named_args, **kwargs): + USE_SPLIT_K = named_args.get("USE_SPLIT_K", kwargs.get("USE_SPLIT_K", None)) + + if USE_SPLIT_K: + pruned_configs = configs + else: + # Deterministic path: a single K block covers the whole reduction, so each program + # stores its result instead of atomic-adding across split-K blocks. + K = named_args.get("K", kwargs.get("K", None)) + block_m = [16, 64, 128] + block_k = align_to(K, 32) + step_k = [256] + warps = [2, 8] + stages = [3, 4] + + pruned_configs = [] + for bm, sk, w, s in itertools.product(block_m, step_k, warps, stages): + pruned_configs.append( + triton.Config( + { + "BLOCK_SIZE_M": bm, + "BLOCK_SIZE_K": block_k, + "STEP_SIZE_K": sk, + }, + num_warps=w, + num_stages=s, + ) + ) + # Triton will skip calling prune function if the autotune returns only one config, which breaks the determinism override here + # So we need to apply NVTE_DISABLE_TRITON_AUTOTUNING in the pruner instead if os.environ.get("NVTE_DISABLE_TRITON_AUTOTUNING", "0") == "1": - configs = configs[:1] - return configs + pruned_configs = pruned_configs[:1] + return pruned_configs -@triton.autotune(configs=projection_config_fwd(), key=["M", "K"], reset_to_zero=["h_ptr", "ms_ptr"]) +@triton.autotune( + configs=projection_config_fwd(), + key=["M", "K", "USE_TMA", "USE_SPLIT_K"], + reset_to_zero=["h_ptr", "ms_ptr"], + prune_configs_by={"early_config_prune": projection_prune_fwd}, +) @triton.jit def _mhc_projection_fwd_fused( x_ptr, # (M, K) @@ -67,12 +95,15 @@ def _mhc_projection_fwd_fused( stride_hm: tl.constexpr, stride_hn: tl.constexpr, stride_ms: tl.constexpr, + stride_norm_weight: tl.constexpr, # Meta-parameters BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, STEP_SIZE_K: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, precision: tl.constexpr, + USE_SPLIT_K: tl.constexpr, # If True, reduce over split-K blocks via atomic_add (non-deterministic); else a single K block stores its result. + USE_TMA: tl.constexpr, # If True, load x and phi via TMA tensor descriptors (Hopper+ only). Falls back to pointer-arith tl.load otherwise. ): pid_m = tl.program_id(axis=0) pid_k = tl.program_id(axis=1) @@ -86,8 +117,9 @@ def _mhc_projection_fwd_fused( tl.assume(stride_hm == 32) tl.assume(stride_hn == 1) tl.assume(stride_ms == 1) + tl.assume(stride_norm_weight == 1) - tl.assume(BLOCK_SIZE_M % 32 == 0) + tl.assume(BLOCK_SIZE_M % 8 == 0) tl.assume(BLOCK_SIZE_K % 32 == 0) tl.assume(BLOCK_SIZE_N == 32) @@ -98,46 +130,101 @@ def _mhc_projection_fwd_fused( h_acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) ms_acc = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32) + if USE_TMA: + x_desc = tl.make_tensor_descriptor( + x_ptr, + shape=[M, K], + strides=[stride_xm, 1], + block_shape=[BLOCK_SIZE_M, STEP_SIZE_K], + ) + phi_desc = tl.make_tensor_descriptor( + phi_ptr, + shape=[N, K], + strides=[stride_phin, 1], + block_shape=[BLOCK_SIZE_N, STEP_SIZE_K], + ) + k_base = pid_k * BLOCK_SIZE_K for k_start in range(0, tl.cdiv(BLOCK_SIZE_K, STEP_SIZE_K)): - k_offs = k_base + k_start * STEP_SIZE_K + tl.arange(0, STEP_SIZE_K) + k_off = k_base + k_start * STEP_SIZE_K + k_offs = k_off + tl.arange(0, STEP_SIZE_K) mask_k = k_offs < K - x_ptrs = x_ptr + offs_m[:, None] * stride_xm + k_offs[None, :] * stride_xk - x = tl.load( - x_ptrs, mask=mask_m[:, None] & mask_k[None, :], other=0.0 - ) # (BLOCK_SIZE_M, BLOCK_SIZE_K) - phi_ptrs = phi_ptr + offs_n_full[:, None] * stride_phin + k_offs[None, :] * stride_phik - phi = tl.load( - phi_ptrs, - mask=(offs_n_full[:, None] < N) & mask_k[None, :], - other=0.0, - cache_modifier=".ca", - ) # (BLOCK_SIZE_N, BLOCK_SIZE_K) - ms_acc += tl.sum(x * x, axis=1) + + if USE_TMA: + x = tl.load_tensor_descriptor(x_desc, [pid_m * BLOCK_SIZE_M, k_off]) + phi = tl.load_tensor_descriptor(phi_desc, [0, k_off]) + else: + x_ptrs = x_ptr + offs_m[:, None] * stride_xm + k_offs[None, :] * stride_xk + x = tl.load( + x_ptrs, mask=mask_m[:, None] & mask_k[None, :], other=0.0 + ) # (BLOCK_SIZE_M, BLOCK_SIZE_K) + phi_ptrs = phi_ptr + offs_n_full[:, None] * stride_phin + k_offs[None, :] * stride_phik + phi = tl.load( + phi_ptrs, + mask=(offs_n_full[:, None] < N) & mask_k[None, :], + other=0.0, + cache_modifier=".ca", + ) # (BLOCK_SIZE_N, BLOCK_SIZE_K) + + ms_acc += tl.sum(x.to(tl.float32) * x.to(tl.float32), axis=1) + + # Currently triton has a bug where for small block size, tl.dot(x, phi.T) will use SMEM to transpose the matrix + # instead of emit a ldmatrix instruction with `.trans` modifier, which leads bank conflicts and performance regression + # See https://github.com/triton-lang/triton/issues/6569#issuecomment-2841739082 h_acc = tl.dot( - x, tl.trans(phi, (1, 0)), h_acc, input_precision=precision, out_dtype=tl.float32 + x.to(phi.dtype), + tl.trans(phi, (1, 0)), + h_acc, + input_precision=precision, + out_dtype=tl.float32, ) h_ptrs = h_ptr + offs_m[:, None] * stride_hm + offs_n_full[None, :] * stride_hn - tl.atomic_add(h_ptrs, h_acc, mask=mask_m[:, None], sem="relaxed") + if USE_SPLIT_K: + tl.atomic_add(h_ptrs, h_acc, mask=mask_m[:, None], sem="relaxed") + else: + tl.store(h_ptrs, h_acc, mask=mask_m[:, None]) offs_ms = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) masks_ms = offs_ms < M offs_ms %= M ms_ptrs = ms_ptr + offs_ms * stride_ms ms = ms_acc / tl.cast(K, tl.float32) - tl.atomic_add(ms_ptrs, ms, mask=masks_ms, sem="relaxed") + if USE_SPLIT_K: + tl.atomic_add(ms_ptrs, ms, mask=masks_ms, sem="relaxed") + else: + tl.store(ms_ptrs, ms, mask=masks_ms) + + +def projection_config_bwd_dx(): + block_m = [32, 128] + block_k = [128] + warps = [2] + stages = [2, 3, 4] + + configs = [] + for m, bk, w, s in itertools.product(block_m, block_k, warps, stages): + configs.append( + triton.Config({"BLOCK_SIZE_M": m, "BLOCK_SIZE_K": bk}, num_warps=w, num_stages=s) + ) + if os.environ.get("NVTE_DISABLE_TRITON_AUTOTUNING", "0") == "1": + configs = configs[:1] + return configs @triton.autotune( - configs=projection_config_bwd(), + configs=projection_config_bwd_dx(), key=["M", "K"], + # When FUSE_GRAD_X_ACC=True the kernel does a read-modify-write on grad_x_ptr; without + # restore_value the autotune timing trials accumulate onto the buffer and corrupt it. + restore_value=["grad_x_ptr"], ) @triton.jit -def _mhc_projection_bwd_fused( +def _mhc_projection_bwd_fused_dx( x_ptr, grad_x_ptr, # (M, K) phi_ptr, # (N, K) + norm_weight_ptr, # (K,) grad_h_ptr, # (M, N) grad_ms_ptr, # (M,) M, @@ -149,6 +236,7 @@ def _mhc_projection_bwd_fused( stride_grad_xk: tl.constexpr, stride_phin, stride_phik: tl.constexpr, + stride_norm_weight: tl.constexpr, stride_grad_phin, stride_grad_phik: tl.constexpr, stride_grad_hm: tl.constexpr, @@ -159,6 +247,8 @@ def _mhc_projection_bwd_fused( BLOCK_SIZE_K: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, precision: tl.constexpr, + FUSE_GRAD_X_ACC: tl.constexpr, + HAS_NORM_WEIGHT: tl.constexpr, ): pid_m = tl.program_id(axis=0) pid_k = tl.program_id(axis=1) @@ -174,6 +264,7 @@ def _mhc_projection_bwd_fused( tl.assume(stride_grad_phin == K) tl.assume(stride_grad_phik == 1) tl.assume(stride_grad_ms == 1) + tl.assume(stride_norm_weight == 1) tl.assume(BLOCK_SIZE_M % 32 == 0) tl.assume(BLOCK_SIZE_K % 32 == 0) @@ -204,19 +295,230 @@ def _mhc_projection_bwd_fused( phi = tl.load( phi_ptrs, mask=(offs_n_full[:, None] < N) & mask_k[None, :], other=0.0 ) # (BLOCK_SIZE_N, BLOCK_SIZE_K) + + if HAS_NORM_WEIGHT: + norm_weight_ptrs = norm_weight_ptr + offs_k * stride_norm_weight + norm_weight = tl.load(norm_weight_ptrs, mask=mask_k, other=0.0, cache_modifier=".ca").to( + phi.dtype + ) # (BLOCK_SIZE_K,) + phi = phi.to(tl.float32) * norm_weight.to(tl.float32)[None, :] + grad_ms = tl.load( grad_ms_ptrs, mask=offs_ms < M, other=0.0, cache_modifier=".ca" ) # (BLOCK_SIZE_M,) grad_x = x * (grad_ms * 2 / tl.cast(K, tl.float32))[:, None] grad_x = tl.dot( - grad_h, phi, acc=grad_x, input_precision=precision, out_dtype=tl.float32 + grad_h.to(phi.dtype), phi, acc=grad_x, input_precision=precision, out_dtype=tl.float32 ) # (BLOCK_SIZE_M, BLOCK_SIZE_K) grad_x_ptrs = grad_x_ptr + offs_m[:, None] * stride_grad_xm + offs_k[None, :] * stride_grad_xk - grad_x = grad_x.to(x.dtype) + if FUSE_GRAD_X_ACC: # If fused gradient accumulation is enabled, the buffer is always fp32 + grad_x_acc = tl.load(grad_x_ptrs, mask=mask_m[:, None] & mask_k[None, :], other=0.0) + grad_x = grad_x.to(tl.float32) + grad_x_acc + else: + grad_x = grad_x.to(x.dtype) tl.store(grad_x_ptrs, grad_x, mask=mask_m[:, None] & mask_k[None, :]) +def projection_config_bwd_dphi(): + block_m = [512, 1024, 2048] + step_m = [32] + block_k = [128, 256] + warps = [2] + stages = [4, 3, 2] + + configs = [] + for bm, sm, bk, w, s in itertools.product(block_m, step_m, block_k, warps, stages): + configs.append( + triton.Config( + {"BLOCK_SIZE_M": bm, "STEP_SIZE_M": sm, "BLOCK_SIZE_K": bk}, + num_warps=w, + num_stages=s, + ) + ) + return configs + + +def projection_prune_bwd_dphi(configs, named_args, **kwargs): + USE_SPLIT_M = named_args.get("USE_SPLIT_M", kwargs.get("USE_SPLIT_M", None)) + M = named_args.get("M", kwargs.get("M", None)) + + if USE_SPLIT_M: + pruned_configs = configs + else: + # Deterministic path: a single M block covers the whole reduction, so each program + # stores grad_phi/grad_norm_weight instead of atomic-adding across split-M blocks. + block_k = [128] + block_m = align_to(M, 128) + step_m = [32] + warps = [4] + # Descending so that when autotuning is disabled (first config taken below), we pick + # the deepest pipeline that fits in shared memory. stages=2 fits on any supported GPU. + stages = [8, 6, 4, 2] + + pruned_configs = [] + for bk, sm, w, s in itertools.product(block_k, step_m, warps, stages): + pruned_configs.append( + triton.Config( + {"BLOCK_SIZE_M": block_m, "STEP_SIZE_M": sm, "BLOCK_SIZE_K": bk}, + num_warps=w, + num_stages=s, + ) + ) + + # Drop configs that exceed the grid Y limit or the device's shared memory + x = named_args.get("x_ptr", kwargs.get("x_ptr", None)) + x_element_size = x.element_size() if x is not None else 4 + max_shared_mem = triton.runtime.driver.active.utils.get_device_properties( + triton.runtime.driver.active.get_current_device() + )["max_shared_mem"] + + def smem_bytes(config): + step_m = config.kwargs["STEP_SIZE_M"] + block_k = config.kwargs["BLOCK_SIZE_K"] + per_stage = step_m * (block_k * x_element_size + 32 * 4) + epilogue = 32 * block_k * 4 + block_k * 4 + return config.num_stages * per_stage + epilogue + + pruned_configs = [ + config + for config in pruned_configs + if triton.cdiv(M, config.kwargs["BLOCK_SIZE_M"]) <= MAX_GRID_DIM_Y + and smem_bytes(config) <= max_shared_mem + ] + + if not pruned_configs: + raise ValueError(f"M={M} exceeds the maximum supported M dimension for this kernel.") + + # Triton will skip calling prune function if the autotune returns only one config, which breaks the determinism override here + # So we need to apply NVTE_DISABLE_TRITON_AUTOTUNING in the pruner instead + if os.environ.get("NVTE_DISABLE_TRITON_AUTOTUNING", "0") == "1": + pruned_configs = pruned_configs[:1] + return pruned_configs + + +@triton.autotune( + configs=projection_config_bwd_dphi(), + key=["M", "K", "USE_SPLIT_M"], + reset_to_zero=["grad_phi_ptr", "grad_norm_weight_ptr"], + prune_configs_by={"early_config_prune": projection_prune_bwd_dphi}, +) +@triton.jit +def _mhc_projection_bwd_fused_dphi( + x_ptr, # (M, K) + grad_H_ptr, # (M, 32) + phi_ptr, # (N, K), N=24 in our case since n = 4 + norm_weight_ptr, # (K,) + grad_phi_ptr, # (N, K), N=24 in our case since n = 4 + grad_norm_weight_ptr, # (K,) + M, + N, + K, + stride_xm, + stride_xk: tl.constexpr, + stride_grad_Hm: tl.constexpr, + stride_grad_Hn: tl.constexpr, + stride_phin, + stride_phik: tl.constexpr, + stride_norm_weight: tl.constexpr, + stride_grad_phin, + stride_grad_phik: tl.constexpr, + stride_grad_norm_weight: tl.constexpr, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + STEP_SIZE_M: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + precision: tl.constexpr, + USE_SPLIT_M: tl.constexpr, # If True, reduce over split-M blocks via atomic_add (non-deterministic); else a single M block stores its result. +): + pid_k = tl.program_id(axis=0) + pid_m = tl.program_id(axis=1) + + tl.assume(pid_k >= 0) + tl.assume(stride_xm > 0) + tl.assume(stride_xk == 1) + tl.assume(stride_grad_Hm == 32) + tl.assume(stride_grad_Hn == 1) + tl.assume(stride_phin == K) + tl.assume(stride_phik == 1) + tl.assume(stride_grad_phin == K) + tl.assume(stride_grad_phin == stride_phin) + tl.assume(stride_grad_phik == 1) + tl.assume(stride_grad_norm_weight == 1) + tl.assume(stride_norm_weight == 1) + + tl.assume(BLOCK_SIZE_M % 128 == 0) + tl.assume(BLOCK_SIZE_K % 64 == 0) + tl.assume(BLOCK_SIZE_N == 32) + tl.assume(STEP_SIZE_M % 32 == 0) + + offs_k = pid_k * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + mask_k = offs_k < K + offs_n_full = tl.arange(0, BLOCK_SIZE_N) + mask_n = offs_n_full < N + + grad_psi_acc = tl.zeros((BLOCK_SIZE_N, BLOCK_SIZE_K), dtype=tl.float32) + + m_start = pid_m * BLOCK_SIZE_M + m_end = tl.minimum(m_start + BLOCK_SIZE_M, M) + for m_idx in range(0, tl.cdiv(m_end - m_start, STEP_SIZE_M)): + offs_m = m_start + m_idx * STEP_SIZE_M + tl.arange(0, STEP_SIZE_M) + mask_m = offs_m < M + x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xk + x = tl.load( + x_ptrs, mask=mask_m[:, None] & mask_k[None, :], other=0.0 + ) # (STEP_SIZE_M, BLOCK_SIZE_K) + grad_H_ptrs = ( + grad_H_ptr + offs_m[:, None] * stride_grad_Hm + offs_n_full[None, :] * stride_grad_Hn + ) + grad_H = tl.load( + grad_H_ptrs, mask=mask_m[:, None] & mask_n[None, :], other=0.0 + ) # (STEP_SIZE_M, BLOCK_SIZE_N) + + grad_psi_acc = tl.dot( + tl.trans(grad_H, (1, 0)), + x.to(grad_H.dtype), + acc=grad_psi_acc, + out_dtype=tl.float32, + input_precision=precision, + ) + + phi_ptrs = phi_ptr + offs_n_full[:, None] * stride_phin + offs_k[None, :] * stride_phik + phi = tl.load( + phi_ptrs, mask=(offs_n_full[:, None] < N) & mask_k[None, :], other=0.0 + ) # (BLOCK_SIZE_N, BLOCK_SIZE_K) + norm_weight_ptrs = norm_weight_ptr + offs_k * stride_norm_weight + norm_weight = tl.load( + norm_weight_ptrs, mask=mask_k, other=0.0, cache_modifier=".cg" + ) # (BLOCK_SIZE_K,) + phi = phi.to(tl.float32) + norm_weight = norm_weight.to(tl.float32) + + # Keep grad_psi in SRAM and get grad_phi & grad_norm_weight + grad_phi = grad_psi_acc * norm_weight[None, :].to(grad_psi_acc.dtype) # (32, BLOCK_SIZE_K) + grad_norm_weight = tl.sum(grad_psi_acc * phi.to(grad_psi_acc.dtype), axis=0) # (BLOCK_SIZE_K,) + + grad_phi_ptrs = ( + grad_phi_ptr + offs_n_full[:, None] * stride_grad_phin + offs_k[None, :] * stride_grad_phik + ) + grad_norm_weight_ptrs = grad_norm_weight_ptr + offs_k * stride_grad_norm_weight + + if USE_SPLIT_M: + tl.atomic_add( + grad_phi_ptrs, + grad_phi, + mask=(offs_n_full[:, None] < N) & mask_k[None, :], + sem="relaxed", + ) + tl.atomic_add(grad_norm_weight_ptrs, grad_norm_weight, mask=mask_k, sem="relaxed") + else: + tl.store( + grad_phi_ptrs, grad_phi.to(phi.dtype), mask=(offs_n_full[:, None] < N) & mask_k[None, :] + ) + tl.store(grad_norm_weight_ptrs, grad_norm_weight.to(norm_weight.dtype), mask=mask_k) + + def scale_config(): block_m = [128] warps = [4] @@ -331,6 +633,8 @@ def _mhc_scale_bwd_fused( grad_b_ptr, # (2n + n^2,) grad_ms_ptr, ms_ptr, # (M,) + ws_grad_a_ptr, # Temporary workspace for a with shape (grid, 4), or None if DETERMINISTIC is False + ws_grad_b_ptr, # Temporary workspace for b with shape (grid, 32), or None if DETERMINISTIC is False M, n, stride_grad_out_m, @@ -349,6 +653,7 @@ def _mhc_scale_bwd_fused( BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, eps: tl.constexpr, + DETERMINISTIC: tl.constexpr, # If True, write per-block partials to a workspace (reduced in the wrapper) instead of atomic_add. ): pid = tl.program_id(0) @@ -416,21 +721,38 @@ def _mhc_scale_bwd_fused( grad_h = tl.where((cols[None, :] >= n) & (cols[None, :] < 2 * n), grad_h_post, grad_h) grad_a = tl.sum(h * grad_h / rms[:, None], axis=0).to(a.dtype) - # Write grad_a[0:4].sum to grad_a_ptr[0], grad_a[4:8].sum to grad_a_ptr[1], and grad_a[8:24].sum to grad_a_ptr[2] - tl.atomic_add(grad_a_ptr, tl.where(cols[None, :] < n, grad_a, 0.0).sum(), sem="relaxed") - tl.atomic_add( - grad_a_ptr + stride_grad_a, - tl.where((cols[None, :] >= n) & (cols[None, :] < 2 * n), grad_a, 0.0).sum(), - sem="relaxed", - ) - tl.atomic_add( - grad_a_ptr + 2 * stride_grad_a, - tl.where((cols[None, :] >= 2 * n) & (cols[None, :] < 2 * n + n * n), grad_a, 0.0).sum(), - sem="relaxed", - ) - grad_b = tl.sum(grad_h, axis=0).to(a.dtype) - tl.atomic_add(grad_b_ptr + cols * stride_grad_b, grad_b, mask=cols < N, sem="relaxed") + + if DETERMINISTIC: + # Deterministic path: each block writes its partials to a workspace row; the wrapper + # sums across blocks. Avoids the non-associative atomic_add reduction over M. + ws_grad_a_ptrs = ws_grad_a_ptr + pid * 4 + # Write grad_a[0:4].sum to ws[pid, 0], grad_a[4:8].sum to ws[pid, 1], grad_a[8:24].sum to ws[pid, 2] + tl.store(ws_grad_a_ptrs, tl.where(cols[None, :] < n, grad_a, 0.0).sum()) + tl.store( + ws_grad_a_ptrs + 1, + tl.where((cols[None, :] >= n) & (cols[None, :] < 2 * n), grad_a, 0.0).sum(), + ) + tl.store( + ws_grad_a_ptrs + 2, + tl.where((cols[None, :] >= 2 * n) & (cols[None, :] < 2 * n + n * n), grad_a, 0.0).sum(), + ) + ws_grad_b_ptrs = ws_grad_b_ptr + pid * 32 + cols + tl.store(ws_grad_b_ptrs, grad_b, mask=cols < N) + else: + # Write grad_a[0:4].sum to grad_a_ptr[0], grad_a[4:8].sum to grad_a_ptr[1], and grad_a[8:24].sum to grad_a_ptr[2] + tl.atomic_add(grad_a_ptr, tl.where(cols[None, :] < n, grad_a, 0.0).sum(), sem="relaxed") + tl.atomic_add( + grad_a_ptr + stride_grad_a, + tl.where((cols[None, :] >= n) & (cols[None, :] < 2 * n), grad_a, 0.0).sum(), + sem="relaxed", + ) + tl.atomic_add( + grad_a_ptr + 2 * stride_grad_a, + tl.where((cols[None, :] >= 2 * n) & (cols[None, :] < 2 * n + n * n), grad_a, 0.0).sum(), + sem="relaxed", + ) + tl.atomic_add(grad_b_ptr + cols * stride_grad_b, grad_b, mask=cols < N, sem="relaxed") grad_rms = (tl.sum((-grad_h * h * a[None, :]), axis=1) / (rms * rms)).to(rms.dtype) grad_ms = grad_rms / (2 * rms) @@ -854,25 +1176,43 @@ def _mhc_sinkhorn_bwd_fused( ) -def aggregate_config(): - block_m = [1, 2, 4] - block_c = [64, 128, 256] - warps = [1, 2, 4] - stages = [1, 2, 3, 4] +def aggregate_config_fwd(): + block_m = [2, 4] + block_c = [256] + warps = [1, 2] + stages = [1, 2, 3] configs = [] for m, c, w, s in itertools.product(block_m, block_c, warps, stages): configs.append( triton.Config({"BLOCK_SIZE_M": m, "BLOCK_SIZE_C": c}, num_warps=w, num_stages=s) ) - if os.environ.get("NVTE_DISABLE_TRITON_AUTOTUNING", "0") == "1": - configs = configs[:1] return configs +def aggregate_prune_fwd(configs, named_args, **kwargs): + M = named_args.get("M", kwargs.get("M", None)) + + pruned_configs = list( + filter( + lambda config: triton.cdiv(M, config.kwargs["BLOCK_SIZE_M"]) <= MAX_GRID_DIM_Y, configs + ) + ) + + if not pruned_configs: + raise ValueError(f"M={M} exceeds the maximum supported M dimension for this kernel.") + + # Triton will skip calling prune function if the autotune returns only one config, which breaks the determinism override here + # So we need to apply NVTE_DISABLE_TRITON_AUTOTUNING in the pruner instead + if os.environ.get("NVTE_DISABLE_TRITON_AUTOTUNING", "0") == "1": + pruned_configs = pruned_configs[:1] + return pruned_configs + + @triton.autotune( - configs=aggregate_config(), + configs=aggregate_config_fwd(), key=["M", "C"], + prune_configs_by={"early_config_prune": aggregate_prune_fwd}, ) @triton.jit def _mhc_aggregate_fwd( @@ -949,7 +1289,66 @@ def _mhc_aggregate_fwd( tl.store(output_ptrs, out, mask=mask_m[:, None] & mask_c[None, :]) -@triton.autotune(configs=aggregate_config(), key=["M", "C"], reset_to_zero=["grad_H_pre_ptr"]) +def aggregate_config_bwd(): + # The real configs are built in `aggregate_prune_bwd` (BLOCK_SIZE_C depends on C at runtime). + # Return a placeholder config so triton won't skip pruning which returns the real configs. + return [ + triton.Config( + {"BLOCK_SIZE_M": 4, "BLOCK_SIZE_C": 256, "STEP_SIZE_C": 64}, num_warps=w, num_stages=2 + ) + for w in (1, 2) + ] + + +def aggregate_prune_bwd(_, named_args, **kwargs): + M = named_args.get("M", kwargs.get("M", None)) + C = named_args.get("C", kwargs.get("C", None)) + block_m = [4] + block_c = align_to(C, 64) + step_c = [64] + warps = [1] + stages = [2, 3, 4] + + pruned_configs = [] + for bm, sc, w, s in itertools.product(block_m, step_c, warps, stages): + pruned_configs.append( + triton.Config( + { + "BLOCK_SIZE_M": bm, + "BLOCK_SIZE_C": block_c, + "STEP_SIZE_C": sc, + }, + num_warps=w, + num_stages=s, + ) + ) + + pruned_configs = list( + filter( + lambda config: triton.cdiv(M, config.kwargs["BLOCK_SIZE_M"]) <= MAX_GRID_DIM_Y, + pruned_configs, + ) + ) + + if not pruned_configs: + raise ValueError(f"M={M} exceeds the maximum supported M dimension for this kernel.") + + # Triton will skip calling prune function if the autotune returns only one config, which breaks the determinism override here + # So we need to apply NVTE_DISABLE_TRITON_AUTOTUNING in the pruner instead + if os.environ.get("NVTE_DISABLE_TRITON_AUTOTUNING", "0") == "1": + pruned_configs = pruned_configs[:1] + return pruned_configs + + +@triton.autotune( + configs=aggregate_config_bwd(), + key=["M", "C"], + reset_to_zero=["grad_H_pre_ptr"], + # When FUSE_GRAD_X_ACC=True the kernel does a read-modify-write on grad_x_ptr; without + # restore_value the autotune timing trials accumulate onto the buffer and corrupt it. + restore_value=["grad_x_ptr"], + prune_configs_by={"early_config_prune": aggregate_prune_bwd}, +) @triton.jit def _mhc_aggregate_bwd( grad_output_ptr, # (M, C) @@ -969,7 +1368,9 @@ def _mhc_aggregate_bwd( # Meta-parameters BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_C: tl.constexpr, + STEP_SIZE_C: tl.constexpr, precision: tl.constexpr, + FUSE_GRAD_X_ACC: tl.constexpr, ): """ Forward: @@ -992,38 +1393,14 @@ def _mhc_aggregate_bwd( tl.assume(stride_grad_output_m > 0 and stride_grad_output_c == 1) tl.assume(BLOCK_SIZE_C % 32 == 0) + tl.assume(STEP_SIZE_C % 32 == 0) + tl.assume(BLOCK_SIZE_C % STEP_SIZE_C == 0) offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_c = pid_c * BLOCK_SIZE_C + tl.arange(0, BLOCK_SIZE_C) - offs_cn = pid_c * BLOCK_SIZE_C * n + tl.arange(0, BLOCK_SIZE_C * n) mask_m = offs_m < M - mask_c = offs_c < C - mask_cn = offs_cn < C * n - - grad_output_ptrs = ( - grad_output_ptr - + offs_m[:, None] * stride_grad_output_m - + offs_c[None, :] * stride_grad_output_c - ) - grad_output = tl.load( - grad_output_ptrs, mask=mask_m[:, None] & mask_c[None, :], other=0.0 - ) # (BLOCK_SIZE_M, BLOCK_SIZE_C) - x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_cn[None, :] * stride_xCn - x = tl.load( - x_ptrs, mask=mask_m[:, None] & mask_cn[None, :], other=0.0 - ) # (BLOCK_SIZE_M, BLOCK_SIZE_C * n) - - grad_H_pre = tl.dot( - tl.reshape(grad_output, (BLOCK_SIZE_M, 1, BLOCK_SIZE_C)), - tl.reshape(x, (BLOCK_SIZE_M, BLOCK_SIZE_C, n)), - input_precision=precision, - out_dtype=tl.float32, - ) - grad_H_pre = tl.reshape(grad_H_pre, (BLOCK_SIZE_M * n,)) # (BLOCK_SIZE_M * n) - offs_grad_H_pre = pid_m * BLOCK_SIZE_M * n + tl.arange(0, BLOCK_SIZE_M * n) - grad_H_pre_ptrs = grad_H_pre_ptr + offs_grad_H_pre - tl.atomic_add(grad_H_pre_ptrs, grad_H_pre, mask=offs_grad_H_pre < M * n, sem="relaxed") + offs_c_start = pid_c * BLOCK_SIZE_C + offs_cn_start = pid_c * BLOCK_SIZE_C * n H_pre_offs = pid_m * BLOCK_SIZE_M * n + tl.arange(0, BLOCK_SIZE_M * n) H_pre = tl.load( @@ -1031,41 +1408,101 @@ def _mhc_aggregate_bwd( ) # (BLOCK_SIZE_M * n) H_pre = tl.reshape(H_pre, (BLOCK_SIZE_M, n)) # (BLOCK_SIZE_M, n) - # grad_x = grad_output @ H_pre.T: (BLOCK_SIZE_M, BLOCK_SIZE_C, 1) @ (BLOCK_SIZE_M, 1, n) = (BLOCK_SIZE_M, BLOCK_SIZE_C, n) - grad_x = grad_output[:, :, None] * H_pre[:, None, :] # (BLOCK_SIZE_M, BLOCK_SIZE_C, n) - grad_x = tl.reshape(grad_x, (BLOCK_SIZE_M, BLOCK_SIZE_C * n)) + grad_H_pre_acc = tl.zeros((BLOCK_SIZE_M, 1, n), dtype=tl.float32) + for i in tl.range(0, BLOCK_SIZE_C, STEP_SIZE_C, loop_unroll_factor=2): + offs_c = offs_c_start + i + tl.arange(0, STEP_SIZE_C) + offs_cn = offs_cn_start + i * n + tl.arange(0, STEP_SIZE_C * n) + mask_c = offs_c < C + mask_cn = offs_cn < C * n + + grad_output_ptrs = ( + grad_output_ptr + + offs_m[:, None] * stride_grad_output_m + + offs_c[None, :] * stride_grad_output_c + ) + grad_output = tl.load( + grad_output_ptrs, mask=mask_m[:, None] & mask_c[None, :], other=0.0 + ) # (BLOCK_SIZE_M, STEP_SIZE_C) - grad_x_ptrs = grad_x_ptr + offs_m[:, None] * stride_grad_xm + offs_cn[None, :] * stride_grad_xCn - tl.store( - grad_x_ptrs, - grad_x, - mask=mask_m[:, None] & mask_cn[None, :], - ) + x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_cn[None, :] * stride_xCn + x = tl.load( + x_ptrs, mask=mask_m[:, None] & mask_cn[None, :], other=0.0 + ) # (BLOCK_SIZE_M, STEP_SIZE_C * n) + + grad_H_pre_acc = tl.dot( + tl.reshape(grad_output, (BLOCK_SIZE_M, 1, STEP_SIZE_C)), + tl.reshape(x, (BLOCK_SIZE_M, STEP_SIZE_C, n)), + acc=grad_H_pre_acc, + input_precision=precision, + out_dtype=tl.float32, + ) + # grad_x = grad_output @ H_pre.T: (BLOCK_SIZE_M, STEP_SIZE_C, 1) @ (BLOCK_SIZE_M, 1, n) = (BLOCK_SIZE_M, STEP_SIZE_C, n) + grad_x = grad_output[:, :, None] * H_pre[:, None, :] # (BLOCK_SIZE_M, STEP_SIZE_C, n) + grad_x = tl.reshape(grad_x, (BLOCK_SIZE_M, STEP_SIZE_C * n)) -def expand_combine_config(): - block_m = [1, 2, 4] - block_c = [128, 256] + grad_x_ptrs = ( + grad_x_ptr + offs_m[:, None] * stride_grad_xm + offs_cn[None, :] * stride_grad_xCn + ) + + if FUSE_GRAD_X_ACC: # If fused gradient accumulation is enabled, the buffer is always fp32 + grad_x_acc = tl.load(grad_x_ptrs, mask=mask_m[:, None] & mask_cn[None, :], other=0.0) + grad_x = grad_x.to(tl.float32) + grad_x_acc + tl.store( + grad_x_ptrs, + grad_x, + mask=mask_m[:, None] & mask_cn[None, :], + ) + + grad_H_pre = tl.reshape(grad_H_pre_acc, (BLOCK_SIZE_M * n,)) # (BLOCK_SIZE_M * n) + offs_grad_H_pre = pid_m * BLOCK_SIZE_M * n + tl.arange(0, BLOCK_SIZE_M * n) + grad_H_pre_ptrs = grad_H_pre_ptr + offs_grad_H_pre + # A single C block covers the reduction, so no atomic_add is needed. + tl.store(grad_H_pre_ptrs, grad_H_pre.to(H_pre.dtype), mask=offs_grad_H_pre < M * n) + + +def expand_combine_config_fwd(): + block_m = [2, 4] + block_c = [256] warps = [1, 2] - stages = [1, 2, 3, 4] + stages = [1, 2, 3] configs = [] for m, c, w, s in itertools.product(block_m, block_c, warps, stages): configs.append( triton.Config({"BLOCK_SIZE_M": m, "BLOCK_SIZE_C": c}, num_warps=w, num_stages=s) ) - if os.environ.get("NVTE_DISABLE_TRITON_AUTOTUNING", "0") == "1": - configs = configs[:1] return configs +def expand_combine_prune_fwd(configs, named_args, **kwargs): + M = named_args.get("M", kwargs.get("M", None)) + + pruned_configs = list( + filter( + lambda config: triton.cdiv(M, config.kwargs["BLOCK_SIZE_M"]) <= MAX_GRID_DIM_Y, configs + ) + ) + + if not pruned_configs: + raise ValueError(f"M={M} exceeds the maximum supported M dimension for this kernel.") + + # Triton will skip calling prune function if the autotune returns only one config, which breaks the determinism override here + # So we need to apply NVTE_DISABLE_TRITON_AUTOTUNING in the pruner instead + if os.environ.get("NVTE_DISABLE_TRITON_AUTOTUNING", "0") == "1": + pruned_configs = pruned_configs[:1] + return pruned_configs + + @triton.autotune( - configs=expand_combine_config(), + configs=expand_combine_config_fwd(), key=["M", "C"], + prune_configs_by={"early_config_prune": expand_combine_prune_fwd}, ) @triton.jit def _mhc_expand_combine_fwd( f_ptr, # (M, C) + bias_ptr, # (C,), or None if HAS_BIAS is False H_post_ptr, # (M, n) x_ptr, # (M, C, n) H_res_ptr, # (M, n, n) @@ -1075,6 +1512,7 @@ def _mhc_expand_combine_fwd( n: tl.constexpr, stride_fm, stride_fc, + stride_bias, # Not used if HAS_BIAS is False stride_xm, stride_xCn, stride_output_m, @@ -1082,9 +1520,10 @@ def _mhc_expand_combine_fwd( # Meta-parameters BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_C: tl.constexpr, + HAS_BIAS: tl.constexpr, ): """ - output = f @ H_post: (BLOCK_SIZE_M, BLOCK_SIZE_C, 1) @ (BLOCK_SIZE_M, 1, n) = (BLOCK_SIZE_M, BLOCK_SIZE_C, n) + output = (f + bias[None, :, None]) @ H_post: (BLOCK_SIZE_M, BLOCK_SIZE_C, 1) @ (BLOCK_SIZE_M, 1, n) = (BLOCK_SIZE_M, BLOCK_SIZE_C, n) + x @ H_res: (BLOCK_SIZE_M, BLOCK_SIZE_C, n) @ (BLOCK_SIZE_M, n, n) = (BLOCK_SIZE_M, BLOCK_SIZE_C, n) """ pid_m = tl.program_id(1) @@ -1095,6 +1534,7 @@ def _mhc_expand_combine_fwd( tl.assume(C > 0) tl.assume(n == 4) tl.assume(stride_fm > 0 and stride_fc == 1) + tl.assume(stride_bias == 1) tl.assume(stride_xm > 0 and stride_xCn == 1) tl.assume(stride_output_m > 0 and stride_output_Cn == 1) @@ -1109,6 +1549,8 @@ def _mhc_expand_combine_fwd( f_ptrs = f_ptr + offs_m[:, None] * stride_fm + offs_c[None, :] * stride_fc f = tl.load(f_ptrs, mask=mask_m[:, None] & mask_c[None, :], other=0.0) + if HAS_BIAS: + bias = tl.load(bias_ptr + offs_c * stride_bias, mask=mask_c, other=0.0) # (BLOCK_SIZE_C,) offs_H_post = pid_m * BLOCK_SIZE_M * n + tl.arange(0, BLOCK_SIZE_M * n) H_post = tl.load( @@ -1116,10 +1558,12 @@ def _mhc_expand_combine_fwd( ) H_post = tl.reshape(H_post, (BLOCK_SIZE_M, n)) # (BLOCK_SIZE_M, n) - # Residual connection path: res_out = f @ H_post: + # Residual connection path: res_out = f @ H_post + bias @ H_post: # (BLOCK_SIZE_M, BLOCK_SIZE_C, 1) @ (BLOCK_SIZE_M, 1, n) = (BLOCK_SIZE_M, n, BLOCK_SIZE_C) # Due to broadcasting, it's equivalent to a multiplicaiton out_acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_C, n), dtype=tl.float32) + if HAS_BIAS: + out_acc = tl.fma(bias[None, :, None], H_post[:, None, :], out_acc) out_acc = tl.fma(f[:, :, None], H_post[:, None, :], out_acc) H_res_offs = pid_m * BLOCK_SIZE_M * n * n + tl.arange(0, BLOCK_SIZE_M * n * n) @@ -1167,332 +1611,75 @@ def _mhc_expand_combine_fwd( tl.store(output_ptrs, out, mask=mask_m[:, None] & mask_cn[None, :]) -@triton.autotune( - configs=expand_combine_config(), - key=["M", "C"], - reset_to_zero=["grad_H_post_ptr", "grad_H_res_ptr"], -) -@triton.jit -def _mhc_expand_combine_bwd( - grad_output_ptr, # (M, C, n) - f_ptr, # (M, C) - H_post_ptr, # (M, n) - x_ptr, # (M, C, n) - H_res_ptr, # (M, n, n) - grad_H_post_ptr, # (M, n) - grad_f_ptr, # (M, C) - grad_H_res_ptr, # (M, n, n) - grad_x_ptr, # (M, C, n) - M, - C, - n: tl.constexpr, - stride_grad_output_m, - stride_grad_output_Cn, - stride_fm, - stride_fc, - stride_xm, - stride_xCn, - stride_grad_fm, - stride_grad_fc, - stride_grad_xm, - stride_grad_xCn, - # Meta-parameters - BLOCK_SIZE_M: tl.constexpr, - BLOCK_SIZE_C: tl.constexpr, - precision: tl.constexpr, -): - """ - Each block - It reads - - (BLOCK_SIZE_M, BLOCK_SIZE_C) of f, which is the output of the attention / FFN module - - (BLOCK_SIZE_M, n) of H_post, which is applied for the transformation of the attention / FFN output - - (BLOCK_SIZE_M, BLOCK_SIZE_C, n) of x, which is the skip connection's input - - (BLOCK_SIZE_M, n*n) of H_res, which is applied for the transformation of the skip connection - and writes - - (BLOCK_SIZE_M, n) of grad_H_post - - (BLOCK_SIZE_M, BLOCK_SIZE_C) of grad_f - - (BLOCK_SIZE_M, n, n) of grad_H_res - - (BLOCK_SIZE_M, BLOCK_SIZE_C, n) of grad_x - - Forward: - out = f @ H_post + x @ H_res - Backward: - GEMM: - grad_H_post = f.T @ grad_output: (BLOCK_SIZE_M, 1, BLOCK_SIZE_C) @ (BLOCK_SIZE_M, BLOCK_SIZE_C, n) = (BLOCK_SIZE_M, 1, n) - grad_H_res = x.T @ grad_output: (BLOCK_SIZE_M, n, BLOCK_SIZE_C) @ (BLOCK_SIZE_M, BLOCK_SIZE_C, n) = (BLOCK_SIZE_M, n, n) - Not GEMM: - grad_f = grad_output @ H_post.T: (BLOCK_SIZE_M, BLOCK_SIZE_C, n) @ (BLOCK_SIZE_M, n, 1) = (BLOCK_SIZE_M, BLOCK_SIZE_C, 1) - grad_x = grad_output @ H_res.T: (BLOCK_SIZE_M, BLOCK_SIZE_C, n) @ (BLOCK_SIZE_M, n, n) = (BLOCK_SIZE_M, BLOCK_SIZE_C, n) - """ - - pid_m = tl.program_id(1) - pid_c = tl.program_id(0) - - tl.static_assert(n == 4) - tl.assume(M > 0) - tl.assume(C > 0) - tl.assume(n == 4) - tl.assume(stride_fm > 0 and stride_fc == 1) - tl.assume(stride_xm > 0 and stride_xCn == 1) - tl.assume(stride_grad_output_m > 0 and stride_grad_output_Cn == 1) - tl.assume(stride_grad_fm > 0 and stride_grad_fc == 1) - tl.assume(stride_grad_xm > 0 and stride_grad_xCn == 1) - - tl.assume(BLOCK_SIZE_C % 32 == 0) - - offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_c = pid_c * BLOCK_SIZE_C + tl.arange(0, BLOCK_SIZE_C) - offs_cn = pid_c * BLOCK_SIZE_C * n + tl.arange(0, BLOCK_SIZE_C * n) - mask_m = offs_m < M - mask_c = offs_c < C - mask_cn = offs_cn < C * n - - f_ptrs = f_ptr + offs_m[:, None] * stride_fm + offs_c[None, :] * stride_fc - f = tl.load(f_ptrs, mask=mask_m[:, None] & mask_c[None, :], other=0.0) - - H_post_offs = pid_m * BLOCK_SIZE_M * n + tl.arange(0, BLOCK_SIZE_M * n) - H_post = tl.load(H_post_ptr + H_post_offs, mask=H_post_offs < M * n, other=0.0) - H_post = tl.reshape(H_post, (BLOCK_SIZE_M, n)) # (BLOCK_SIZE_M, n) - - H_res_offs = pid_m * BLOCK_SIZE_M * n * n + tl.arange(0, BLOCK_SIZE_M * n * n) - H_res = tl.load( - H_res_ptr + H_res_offs, mask=H_res_offs < M * n * n, other=0.0 - ) # (BLOCK_SIZE_M, n, n) - H_res = tl.reshape(H_res, (BLOCK_SIZE_M, n, n)) # (BLOCK_SIZE_M, n, n) - - grad_out_ptrs = ( - grad_output_ptr - + offs_m[:, None] * stride_grad_output_m - + offs_cn[None, :] * stride_grad_output_Cn - ) - grad_out = tl.load( - grad_out_ptrs, mask=mask_m[:, None] & mask_cn[None, :], other=0.0 - ) # (BLOCK_SIZE_M, BLOCK_SIZE_C * n) - grad_out = tl.reshape( - grad_out, (BLOCK_SIZE_M, BLOCK_SIZE_C, n) - ) # (BLOCK_SIZE_M, BLOCK_SIZE_C, n) - - # grad_H_post = f.T @ grad_output # (BLOCK_SIZE_M, 1, BLOCK_SIZE_C) @ (BLOCK_SIZE_M, BLOCK_SIZE_C, n) = (BLOCK_SIZE_M, 1, n) - grad_H_post = tl.dot( - tl.reshape(f, (BLOCK_SIZE_M, 1, BLOCK_SIZE_C)), - tl.reshape(grad_out, (BLOCK_SIZE_M, BLOCK_SIZE_C, n)), - input_precision=precision, - out_dtype=tl.float32, - ) # (BLOCK_SIZE_M, 1, n) - grad_H_post = tl.reshape(grad_H_post, (BLOCK_SIZE_M * n,)) # (BLOCK_SIZE_M * n) - offs_grad_H_post = pid_m * BLOCK_SIZE_M * n + tl.arange(0, BLOCK_SIZE_M * n) - grad_H_post_ptrs = grad_H_post_ptr + offs_grad_H_post - tl.atomic_add(grad_H_post_ptrs, grad_H_post, mask=offs_grad_H_post < M * n, sem="relaxed") - - x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_cn[None, :] * stride_xCn - x = tl.load( - x_ptrs, mask=mask_m[:, None] & mask_cn[None, :], other=0.0 - ) # (BLOCK_SIZE_M, BLOCK_SIZE_C*n) - x = tl.reshape(x, (BLOCK_SIZE_M, BLOCK_SIZE_C, n)) # (BLOCK_SIZE_M, BLOCK_SIZE_C, n) - - # grad_H_res = x.T @ grad_output: (BLOCK_SIZE_M, n, BLOCK_SIZE_C) @ (BLOCK_SIZE_M, BLOCK_SIZE_C, n) = (BLOCK_SIZE_M, n, n) - grad_H_res = tl.dot( - tl.trans(x, (0, 2, 1)), grad_out, input_precision=precision, out_dtype=tl.float32 - ) # (BLOCK_SIZE_M, n, n) - grad_H_res = tl.reshape(grad_H_res, (BLOCK_SIZE_M * n * n,)) # (BLOCK_SIZE_M * n * n) - offs_grad_H_res = pid_m * BLOCK_SIZE_M * n * n + tl.arange(0, BLOCK_SIZE_M * n * n) - grad_H_res_ptrs = grad_H_res_ptr + offs_grad_H_res - tl.atomic_add( - grad_H_res_ptrs, grad_H_res.to(tl.float32), mask=offs_grad_H_res < M * n * n, sem="relaxed" - ) - - grad_out_reshape = tl.reshape( - grad_out, (BLOCK_SIZE_M, BLOCK_SIZE_C, 2, 2) - ) # (BLOCK_SIZE_M, BLOCK_SIZE_C, 2, 2) - grad_out01, grad_out23 = tl.split( - grad_out_reshape - ) # (BLOCK_SIZE_M, BLOCK_SIZE_C, 2), (BLOCK_SIZE_M, BLOCK_SIZE_C, 2) - grad_out0, grad_out1 = tl.split( - grad_out01 - ) # (BLOCK_SIZE_M, BLOCK_SIZE_C), (BLOCK_SIZE_M, BLOCK_SIZE_C) - grad_out2, grad_out3 = tl.split( - grad_out23 - ) # (BLOCK_SIZE_M, BLOCK_SIZE_C), (BLOCK_SIZE_M, BLOCK_SIZE_C) - - # grad_f = grad_output @ H_post.T: (BLOCK_SIZE_M, 1, n) @ (BLOCK_SIZE_M, n, BLOCK_SIZE_C) = (BLOCK_SIZE_M, 1, BLOCK_SIZE_C) - # Triton doesn't support dot prod with inner dimension < 16, so we need to hack this: - # grad_f = grad_out[:, :, 0] @ H_post.T[:, 0, :] (BLOCK_SIZE_M, BLOCK_SIZE_C, 1) @ (BLOCK_SIZE_M, 1, 1) - # + grad_out[:, :, 1] @ H_post.T[:, 1, :] - # + grad_out[:, :, 2] @ H_post.T[:, 2, :] - # + grad_out[:, :, 3] @ H_post.T[:, 3, :] - # where H_post.T[:, i, :] = H_post[:, :, i] - H_post = tl.reshape(H_post, (BLOCK_SIZE_M, 2, 2)) - H_post01, H_post23 = tl.split(H_post) # (BLOCK_SIZE_M, 2), (BLOCK_SIZE_M, 2) - H_post0, H_post1 = tl.split(H_post01) # (BLOCK_SIZE_M,), (BLOCK_SIZE_M,) - H_post2, H_post3 = tl.split(H_post23) # (BLOCK_SIZE_M,), (BLOCK_SIZE_M,) - - grad_f_acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_C), dtype=tl.float32) - # (BLOCK_SIZE_M, BLOCK_SIZE_C) * (BLOCK_SIZE_M, 1) -> (BLOCK_SIZE_M, BLOCK_SIZE_C) - grad_f_acc = tl.fma(grad_out0, H_post0[:, None], grad_f_acc) - grad_f_acc = tl.fma(grad_out1, H_post1[:, None], grad_f_acc) - grad_f_acc = tl.fma(grad_out2, H_post2[:, None], grad_f_acc) - grad_f_acc = tl.fma(grad_out3, H_post3[:, None], grad_f_acc) - grad_f = grad_f_acc.to(f.dtype) - - grad_f_ptrs = grad_f_ptr + offs_m[:, None] * stride_grad_fm + offs_c[None, :] * stride_grad_fc - tl.store(grad_f_ptrs, grad_f, mask=mask_m[:, None] & mask_c[None, :]) - - # grad_x = grad_output @ H_res.T: (BLOCK_SIZE_M, BLOCK_SIZE_C, n) @ (BLOCK_SIZE_M, n, n) = (BLOCK_SIZE_M, n, BLOCK_SIZE_C) - # The inner dim is n=4 which is too small for triton, so we will manually unroll the matmul - # grad_x = grad_out[:, :, 0] @ H_res.T[:, 0, :] - # + grad_out[:, :, 1] @ H_res.T[:, 1, :] - # + grad_out[:, :, 2] @ H_res.T[:, 2, :] - # + grad_out[:, :, 3] @ H_res.T[:, 3, :] - # where H_res.T[:, i, :] = H_res[:, :, i] - # Due to broadcasting, it's equivalent to multiplying each H_res[:, i, :].T with grad_out[:, i, :] - - H_res_reshape = tl.reshape(H_res, (BLOCK_SIZE_M, n, 2, 2)) # (BLOCK_SIZE_M, n, 2, 2) - H_res01, H_res23 = tl.split(H_res_reshape) # (BLOCK_SIZE_M, n, 2), (BLOCK_SIZE_M, n, 2) - H_res0, H_res1 = tl.split(H_res01) # (BLOCK_SIZE_M, n), (BLOCK_SIZE_M, n) - H_res2, H_res3 = tl.split(H_res23) # (BLOCK_SIZE_M, n), (BLOCK_SIZE_M, n) - - grad_x_acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_C, n), dtype=tl.float32) - grad_x_acc = tl.fma(grad_out0[:, :, None], H_res0[:, None, :], grad_x_acc) - grad_x_acc = tl.fma(grad_out1[:, :, None], H_res1[:, None, :], grad_x_acc) - grad_x_acc = tl.fma(grad_out2[:, :, None], H_res2[:, None, :], grad_x_acc) - grad_x_acc = tl.fma(grad_out3[:, :, None], H_res3[:, None, :], grad_x_acc) - - grad_x = grad_x_acc.to(x.dtype) - grad_x = tl.reshape(grad_x, (BLOCK_SIZE_M, BLOCK_SIZE_C * n)) # (BLOCK_SIZE_M, BLOCK_SIZE_C*n) - - grad_x_ptrs = grad_x_ptr + offs_m[:, None] * stride_grad_xm + offs_cn[None, :] * stride_grad_xCn - tl.store(grad_x_ptrs, grad_x, mask=mask_m[:, None] & mask_cn[None, :]) - - -@triton.autotune( - configs=expand_combine_config(), - key=["M", "C"], -) -@triton.jit -def _mhc_expand_combine_with_bias_fwd( - f_ptr, # (M, C) - bias_ptr, # (C,) - H_post_ptr, # (M, n) - x_ptr, # (M, C, n) - H_res_ptr, # (M, n, n) - output_ptr, # # (M, C, n) - M, - C, - n: tl.constexpr, - stride_fm, - stride_fc, - stride_bias, - stride_xm, - stride_xCn, - stride_output_m, - stride_output_Cn, - # Meta-parameters - BLOCK_SIZE_M: tl.constexpr, - BLOCK_SIZE_C: tl.constexpr, -): - """ - output = (f + bias[None, :, None]) @ H_post: (BLOCK_SIZE_M, BLOCK_SIZE_C, 1) @ (BLOCK_SIZE_M, 1, n) = (BLOCK_SIZE_M, BLOCK_SIZE_C, n) - + x @ H_res: (BLOCK_SIZE_M, BLOCK_SIZE_C, n) @ (BLOCK_SIZE_M, n, n) = (BLOCK_SIZE_M, BLOCK_SIZE_C, n) - """ - pid_m = tl.program_id(1) - pid_c = tl.program_id(0) - - tl.static_assert(n == 4) - tl.assume(M > 0) - tl.assume(C > 0) - tl.assume(n == 4) - tl.assume(stride_fm > 0 and stride_fc == 1) - tl.assume(stride_bias == 1) - tl.assume(stride_xm > 0 and stride_xCn == 1) - tl.assume(stride_output_m > 0 and stride_output_Cn == 1) - - tl.assume(BLOCK_SIZE_C % 32 == 0) - - offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_c = pid_c * BLOCK_SIZE_C + tl.arange(0, BLOCK_SIZE_C) - offs_cn = pid_c * BLOCK_SIZE_C * n + tl.arange(0, BLOCK_SIZE_C * n) - mask_m = offs_m < M - mask_c = offs_c < C - mask_cn = offs_cn < C * n +def expand_combine_config_bwd(): + # The real configs are built in `expand_combine_prune_bwd` (BLOCK_SIZE_C depends on C at runtime). + # Return a placeholder config so triton won't skip pruning which returns the real configs. + return [ + triton.Config( + {"BLOCK_SIZE_M": 4, "BLOCK_SIZE_C": 256, "STEP_SIZE_C": 64}, num_warps=w, num_stages=2 + ) + for w in (1, 2) + ] - f_ptrs = f_ptr + offs_m[:, None] * stride_fm + offs_c[None, :] * stride_fc - f = tl.load(f_ptrs, mask=mask_m[:, None] & mask_c[None, :], other=0.0) - bias = tl.load(bias_ptr + offs_c * stride_bias, mask=mask_c, other=0.0) # (BLOCK_SIZE_C,) - offs_H_post = pid_m * BLOCK_SIZE_M * n + tl.arange(0, BLOCK_SIZE_M * n) - H_post = tl.load( - H_post_ptr + offs_H_post, mask=offs_H_post < M * n, other=0.0, cache_modifier=".ca" - ) - H_post = tl.reshape(H_post, (BLOCK_SIZE_M, n)) # (BLOCK_SIZE_M, n) +def expand_combine_prune_bwd(_, named_args, **kwargs): + M = named_args.get("M", kwargs.get("M", None)) + C = named_args.get("C", kwargs.get("C", None)) + block_m = [4] + block_c = align_to(C, 32) + step_c = [64, 128] + warps = [1, 2] + stages = [2, 3] - # Residual connection path: res_out = f @ H_post + bias @ H_post: - # (BLOCK_SIZE_M, BLOCK_SIZE_C, 1) @ (BLOCK_SIZE_M, 1, n) = (BLOCK_SIZE_M, n, BLOCK_SIZE_C) - # Due to broadcasting, it's equivalent to a multiplicaiton - out_acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_C, n), dtype=tl.float32) - out_acc = tl.fma(bias[None, :, None], H_post[:, None, :], out_acc) - out_acc = tl.fma(f[:, :, None], H_post[:, None, :], out_acc) + pruned_configs = [] + for bm, sc, w, s in itertools.product(block_m, step_c, warps, stages): + pruned_configs.append( + triton.Config( + { + "BLOCK_SIZE_M": bm, + "BLOCK_SIZE_C": block_c, + "STEP_SIZE_C": sc, + }, + num_warps=w, + num_stages=s, + ) + ) - H_res_offs = pid_m * BLOCK_SIZE_M * n * n + tl.arange(0, BLOCK_SIZE_M * n * n) - H_res = tl.load( - H_res_ptr + H_res_offs, mask=H_res_offs < M * n * n, other=0.0, cache_modifier=".ca" + pruned_configs = list( + filter( + lambda config: triton.cdiv(M, config.kwargs["BLOCK_SIZE_M"]) <= MAX_GRID_DIM_Y, + pruned_configs, + ) ) - H_res = tl.reshape(H_res, (BLOCK_SIZE_M, n, n)) # (BLOCK_SIZE_M, n, n) - x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_cn[None, :] * stride_xCn - x = tl.load( - x_ptrs, mask=mask_m[:, None] & mask_cn[None, :], other=0.0 - ) # (BLOCK_SIZE_M, BLOCK_SIZE_C, n) - - # Manifold connection path: manifold_out = H_res @ x: - # (BLOCK_SIZE_M, BLOCK_SIZE_C, n) @ (BLOCK_SIZE_M, n, n) = (BLOCK_SIZE_M, BLOCK_SIZE_C, n) - # triton doesn't support dot prod with inner dimension < 16, so we need to manually unroll the computation for n=4: - # x @ H_res = x[:, :, 0] @ H_res[:, 0, :] - # + x[:, :, 1] @ H_res[:, 1, :] - # + x[:, :, 2] @ H_res[:, 2, :] - # + x[:, :, 3] @ H_res[:, 3, :] - - x_reshape = tl.reshape(x, (BLOCK_SIZE_M, BLOCK_SIZE_C, 2, 2)) - x01, x23 = tl.split( - x_reshape - ) # (BLOCK_SIZE_M, BLOCK_SIZE_C, 2), (BLOCK_SIZE_M, BLOCK_SIZE_C, 2) - x0, x1 = tl.split(x01) # (BLOCK_SIZE_M, BLOCK_SIZE_C), (BLOCK_SIZE_M, BLOCK_SIZE_C) - x2, x3 = tl.split(x23) # (BLOCK_SIZE_M, BLOCK_SIZE_C), (BLOCK_SIZE_M, BLOCK_SIZE_C) - - H_resT = tl.reshape(tl.trans(H_res, (0, 2, 1)), (BLOCK_SIZE_M, n, 2, 2)) - H_res01, H_res23 = tl.split(H_resT) # (BLOCK_SIZE_M, n, 2), (BLOCK_SIZE_M, n, 2) - H_res0, H_res1 = tl.split(H_res01) # (BLOCK_SIZE_M, n), (BLOCK_SIZE_M, n) - H_res2, H_res3 = tl.split(H_res23) # (BLOCK_SIZE_M, n), (BLOCK_SIZE_M, n) + if not pruned_configs: + raise ValueError(f"M={M} exceeds the maximum supported M dimension for this kernel.") - out_acc = tl.fma(x0[:, :, None], H_res0[:, None, :], out_acc) - out_acc = tl.fma(x1[:, :, None], H_res1[:, None, :], out_acc) - out_acc = tl.fma(x2[:, :, None], H_res2[:, None, :], out_acc) - out_acc = tl.fma(x3[:, :, None], H_res3[:, None, :], out_acc) - - out = out_acc.to(x.dtype) - out = tl.reshape(out, (BLOCK_SIZE_M, BLOCK_SIZE_C * n)) # (BLOCK_SIZE_M, BLOCK_SIZE_C*n) - - output_ptrs = ( - output_ptr + offs_m[:, None] * stride_output_m + offs_cn[None, :] * stride_output_Cn - ) - tl.store(output_ptrs, out, mask=mask_m[:, None] & mask_cn[None, :]) + # Triton will skip calling prune function if the autotune returns only one config, which breaks the determinism override here + # So we need to apply NVTE_DISABLE_TRITON_AUTOTUNING in the pruner instead + if os.environ.get("NVTE_DISABLE_TRITON_AUTOTUNING", "0") == "1": + pruned_configs = pruned_configs[:1] + return pruned_configs @triton.autotune( - configs=expand_combine_config(), - key=["M", "C"], + configs=expand_combine_config_bwd(), + key=["M", "C", "DETERMINISTIC"], reset_to_zero=["grad_H_post_ptr", "grad_H_res_ptr", "grad_bias_ptr"], + prune_configs_by={"early_config_prune": expand_combine_prune_bwd}, ) @triton.jit -def _mhc_expand_combine_with_bias_bwd( +def _mhc_expand_combine_bwd( grad_output_ptr, # (M, C, n) f_ptr, # (M, C) - bias_ptr, # (C,) + bias_ptr, # (C,), or None if HAS_BIAS is False H_post_ptr, # (M, n) x_ptr, # (M, C, n) H_res_ptr, # (M, n, n) grad_H_post_ptr, # (M, n) grad_f_ptr, # (M, C) - grad_bias_ptr, # (C,) + grad_bias_ptr, # (C,), or None if HAS_BIAS is False + grad_bias_ws_ptr, # (grid_m, C), or None if HAS_BIAS is False or DETERMINISTIC is False grad_H_res_ptr, # (M, n, n) grad_x_ptr, # (M, C, n) M, @@ -1502,18 +1689,24 @@ def _mhc_expand_combine_with_bias_bwd( stride_grad_output_Cn, stride_fm, stride_fc, - stride_bias, + stride_bias, # Not used if HAS_BIAS is False stride_xm, stride_xCn, stride_grad_fm, stride_grad_fc, - stride_grad_bias, + stride_grad_bias, # Not used if HAS_BIAS is False + stride_grad_bias_ws_m, # Not used if HAS_BIAS is False or DETERMINISTIC is False + stride_grad_bias_ws_c, # Not used if HAS_BIAS is False or DETERMINISTIC is False stride_grad_xm, stride_grad_xCn, # Meta-parameters BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_C: tl.constexpr, + STEP_SIZE_C: tl.constexpr, precision: tl.constexpr, + HAS_BIAS: tl.constexpr, + FUSE_GRAD_X_ACC: tl.constexpr, + DETERMINISTIC: tl.constexpr, # If True, grad_bias partials go to a workspace (reduced in the wrapper) instead of atomic_add. ): """ Each block @@ -1557,137 +1750,173 @@ def _mhc_expand_combine_with_bias_bwd( tl.assume(BLOCK_SIZE_C % 32 == 0) offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_c = pid_c * BLOCK_SIZE_C + tl.arange(0, BLOCK_SIZE_C) - offs_cn = pid_c * BLOCK_SIZE_C * n + tl.arange(0, BLOCK_SIZE_C * n) mask_m = offs_m < M - mask_c = offs_c < C - mask_cn = offs_cn < C * n - f_ptrs = f_ptr + offs_m[:, None] * stride_fm + offs_c[None, :] * stride_fc - f = tl.load(f_ptrs, mask=mask_m[:, None] & mask_c[None, :], other=0.0) + offs_c_start = pid_c * BLOCK_SIZE_C + offs_cn_start = pid_c * BLOCK_SIZE_C * n - bias = tl.load(bias_ptr + offs_c * stride_bias, mask=mask_c, other=0.0) # (BLOCK_SIZE_C,) + grad_H_post_acc = tl.zeros((BLOCK_SIZE_M, 1, n), dtype=tl.float32) + grad_H_res_acc = tl.zeros((BLOCK_SIZE_M, n, n), dtype=tl.float32) H_post_offs = pid_m * BLOCK_SIZE_M * n + tl.arange(0, BLOCK_SIZE_M * n) H_post = tl.load(H_post_ptr + H_post_offs, mask=H_post_offs < M * n, other=0.0) - H_post = tl.reshape(H_post, (BLOCK_SIZE_M, n)) # (BLOCK_SIZE_M, n) + H_post_reshape = tl.reshape(H_post, (BLOCK_SIZE_M, 2, 2)) + H_post01, H_post23 = tl.split(H_post_reshape) # (BLOCK_SIZE_M, 2), (BLOCK_SIZE_M, 2) + H_post0, H_post1 = tl.split(H_post01) # (BLOCK_SIZE_M,), (BLOCK_SIZE_M,) + H_post2, H_post3 = tl.split(H_post23) # (BLOCK_SIZE_M,), (BLOCK_SIZE_M,) H_res_offs = pid_m * BLOCK_SIZE_M * n * n + tl.arange(0, BLOCK_SIZE_M * n * n) H_res = tl.load( H_res_ptr + H_res_offs, mask=H_res_offs < M * n * n, other=0.0 ) # (BLOCK_SIZE_M, n, n) H_res = tl.reshape(H_res, (BLOCK_SIZE_M, n, n)) # (BLOCK_SIZE_M, n, n) + H_res_reshape = tl.reshape(H_res, (BLOCK_SIZE_M, n, 2, 2)) # (BLOCK_SIZE_M, n, 2, 2) + H_res01, H_res23 = tl.split(H_res_reshape) # (BLOCK_SIZE_M, n, 2), (BLOCK_SIZE_M, n, 2) + H_res0, H_res1 = tl.split(H_res01) # (BLOCK_SIZE_M, n), (BLOCK_SIZE_M, n) + H_res2, H_res3 = tl.split(H_res23) # (BLOCK_SIZE_M, n), (BLOCK_SIZE_M, n) - grad_out_ptrs = ( - grad_output_ptr - + offs_m[:, None] * stride_grad_output_m - + offs_cn[None, :] * stride_grad_output_Cn - ) - grad_out = tl.load( - grad_out_ptrs, mask=mask_m[:, None] & mask_cn[None, :], other=0.0 - ) # (BLOCK_SIZE_M, BLOCK_SIZE_C * n) - grad_out = tl.reshape( - grad_out, (BLOCK_SIZE_M, BLOCK_SIZE_C, n) - ) # (BLOCK_SIZE_M, BLOCK_SIZE_C, n) + for i in tl.range(0, BLOCK_SIZE_C, STEP_SIZE_C, loop_unroll_factor=2): + offs_c = offs_c_start + i + tl.arange(0, STEP_SIZE_C) + offs_cn = offs_cn_start + i * n + tl.arange(0, STEP_SIZE_C * n) + mask_c = offs_c < C + mask_cn = offs_cn < C * n - # grad_H_post = f.T @ grad_output # (BLOCK_SIZE_M, 1, BLOCK_SIZE_C) @ (BLOCK_SIZE_M, BLOCK_SIZE_C, n) = (BLOCK_SIZE_M, 1, n) - grad_H_post = tl.dot( - tl.reshape(f, (BLOCK_SIZE_M, 1, BLOCK_SIZE_C)), - tl.reshape(grad_out, (BLOCK_SIZE_M, BLOCK_SIZE_C, n)), - input_precision=precision, - out_dtype=tl.float32, - ) # (BLOCK_SIZE_M, 1, n) - grad_H_post = tl.dot( - tl.broadcast_to(bias[None, None, :], (BLOCK_SIZE_M, 1, BLOCK_SIZE_C)), - tl.reshape(grad_out, (BLOCK_SIZE_M, BLOCK_SIZE_C, n)), - acc=grad_H_post, - input_precision=precision, - out_dtype=tl.float32, - ) # (BLOCK_SIZE_M, 1, n) - grad_H_post = tl.reshape(grad_H_post, (BLOCK_SIZE_M * n,)) # (BLOCK_SIZE_M * n) + f_ptrs = f_ptr + offs_m[:, None] * stride_fm + offs_c[None, :] * stride_fc + f = tl.load(f_ptrs, mask=mask_m[:, None] & mask_c[None, :], other=0.0) + + if HAS_BIAS: + bias = tl.load( + bias_ptr + offs_c * stride_bias, mask=mask_c, other=0.0 + ) # (STEP_SIZE_C,) + + grad_out_ptrs = ( + grad_output_ptr + + offs_m[:, None] * stride_grad_output_m + + offs_cn[None, :] * stride_grad_output_Cn + ) + grad_out = tl.load( + grad_out_ptrs, mask=mask_m[:, None] & mask_cn[None, :], other=0.0 + ) # (BLOCK_SIZE_M, STEP_SIZE_C * n) + grad_out = tl.reshape( + grad_out, (BLOCK_SIZE_M, STEP_SIZE_C, n) + ) # (BLOCK_SIZE_M, STEP_SIZE_C, n) + + # grad_H_post = f.T @ grad_output # (BLOCK_SIZE_M, 1, STEP_SIZE_C) @ (BLOCK_SIZE_M, STEP_SIZE_C, n) = (BLOCK_SIZE_M, 1, n) + grad_H_post_acc = tl.dot( + tl.reshape(f, (BLOCK_SIZE_M, 1, STEP_SIZE_C)), + tl.reshape(grad_out, (BLOCK_SIZE_M, STEP_SIZE_C, n)), + acc=grad_H_post_acc, + input_precision=precision, + out_dtype=tl.float32, + ) # (BLOCK_SIZE_M, 1, n) + if HAS_BIAS: + grad_H_post_acc = tl.dot( + tl.broadcast_to(bias[None, None, :], (BLOCK_SIZE_M, 1, STEP_SIZE_C)), + tl.reshape(grad_out, (BLOCK_SIZE_M, STEP_SIZE_C, n)), + acc=grad_H_post_acc, + input_precision=precision, + out_dtype=tl.float32, + ) # (BLOCK_SIZE_M, 1, n) + + x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_cn[None, :] * stride_xCn + x = tl.load( + x_ptrs, mask=mask_m[:, None] & mask_cn[None, :], other=0.0 + ) # (BLOCK_SIZE_M, STEP_SIZE_C*n) + x = tl.reshape(x, (BLOCK_SIZE_M, STEP_SIZE_C, n)) # (BLOCK_SIZE_M, STEP_SIZE_C, n) + + # grad_H_res = x.T @ grad_output: (BLOCK_SIZE_M, n, STEP_SIZE_C) @ (BLOCK_SIZE_M, STEP_SIZE_C, n) = (BLOCK_SIZE_M, n, n) + grad_H_res_acc = tl.dot( + tl.trans(x, (0, 2, 1)), + grad_out, + acc=grad_H_res_acc, + input_precision=precision, + out_dtype=tl.float32, + ) # (BLOCK_SIZE_M, n, n) + + grad_out_reshape = tl.reshape( + grad_out, (BLOCK_SIZE_M, STEP_SIZE_C, 2, 2) + ) # (BLOCK_SIZE_M, STEP_SIZE_C, 2, 2) + grad_out01, grad_out23 = tl.split( + grad_out_reshape + ) # (BLOCK_SIZE_M, STEP_SIZE_C, 2), (BLOCK_SIZE_M, STEP_SIZE_C, 2) + grad_out0, grad_out1 = tl.split( + grad_out01 + ) # (BLOCK_SIZE_M, STEP_SIZE_C), (BLOCK_SIZE_M, STEP_SIZE_C) + grad_out2, grad_out3 = tl.split( + grad_out23 + ) # (BLOCK_SIZE_M, STEP_SIZE_C), (BLOCK_SIZE_M, STEP_SIZE_C) + + # grad_f = grad_output @ H_post.T: (BLOCK_SIZE_M, 1, n) @ (BLOCK_SIZE_M, n, STEP_SIZE_C) = (BLOCK_SIZE_M, 1, STEP_SIZE_C) + # Triton doesn't support dot prod with inner dimension < 16, so we need to hack this: + # = grad_out[:, :, 0] @ H_post.T[:, 0, :] (BLOCK_SIZE_M, STEP_SIZE_C, 1) @ (BLOCK_SIZE_M, 1, 1) + # + grad_out[:, :, 1] @ H_post.T[:, 1, :] + # + grad_out[:, :, 2] @ H_post.T[:, 2, :] + # + grad_out[:, :, 3] @ H_post.T[:, 3, :] + # where H_post.T[:, i, :] = H_post[:, :, i] + + grad_f_acc = tl.zeros((BLOCK_SIZE_M, STEP_SIZE_C), dtype=tl.float32) + # (BLOCK_SIZE_M, STEP_SIZE_C) * (BLOCK_SIZE_M, 1) -> (BLOCK_SIZE_M, STEP_SIZE_C) + grad_f_acc = tl.fma(grad_out0, H_post0[:, None], grad_f_acc) + grad_f_acc = tl.fma(grad_out1, H_post1[:, None], grad_f_acc) + grad_f_acc = tl.fma(grad_out2, H_post2[:, None], grad_f_acc) + grad_f_acc = tl.fma(grad_out3, H_post3[:, None], grad_f_acc) + grad_f = grad_f_acc.to(f.dtype) + + grad_f_ptrs = ( + grad_f_ptr + offs_m[:, None] * stride_grad_fm + offs_c[None, :] * stride_grad_fc + ) + tl.store(grad_f_ptrs, grad_f, mask=mask_m[:, None] & mask_c[None, :]) + + if HAS_BIAS: + grad_bias = tl.sum(grad_f_acc, axis=0) # (STEP_SIZE_C,) + # This is reduction over M dimension, so it has nothing to do with whether we use split-C. It only depends on determinism or not. + if DETERMINISTIC: + grad_bias_ws_ptrs = ( + grad_bias_ws_ptr + + pid_m * stride_grad_bias_ws_m + + offs_c * stride_grad_bias_ws_c + ) + tl.store(grad_bias_ws_ptrs, grad_bias, mask=mask_c) + else: + grad_bias_ptrs = grad_bias_ptr + offs_c * stride_grad_bias + tl.atomic_add(grad_bias_ptrs, grad_bias, mask=mask_c, sem="relaxed") + + # grad_x = grad_output @ H_res.T: (BLOCK_SIZE_M, STEP_SIZE_C, n) @ (BLOCK_SIZE_M, n, n) = (BLOCK_SIZE_M, n, STEP_SIZE_C) + # The inner dim is n=4 which is too small for triton, so we will manually unroll the matmul + # grad_x = grad_out[:, :, 0] @ H_res.T[:, 0, :] + # + grad_out[:, :, 1] @ H_res.T[:, 1, :] + # + grad_out[:, :, 2] @ H_res.T[:, 2, :] + # + grad_out[:, :, 3] @ H_res.T[:, 3, :] + # where H_res.T[:, i, :] = H_res[:, :, i] + # Due to broadcasting, it's equivalent to multiplying each H_res[:, i, :].T with grad_out[:, i, :] + + grad_x_acc = tl.zeros((BLOCK_SIZE_M, STEP_SIZE_C, n), dtype=tl.float32) + grad_x_acc = tl.fma(grad_out0[:, :, None], H_res0[:, None, :], grad_x_acc) + grad_x_acc = tl.fma(grad_out1[:, :, None], H_res1[:, None, :], grad_x_acc) + grad_x_acc = tl.fma(grad_out2[:, :, None], H_res2[:, None, :], grad_x_acc) + grad_x_acc = tl.fma(grad_out3[:, :, None], H_res3[:, None, :], grad_x_acc) + + if FUSE_GRAD_X_ACC: + grad_x = grad_x_acc # If fusing gradient accumulation, the buffer should be always fp32 so we don't cast here + else: + grad_x = grad_x_acc.to(x.dtype) + grad_x = tl.reshape( + grad_x, (BLOCK_SIZE_M, STEP_SIZE_C * n) + ) # (BLOCK_SIZE_M, STEP_SIZE_C*n) + + grad_x_ptrs = ( + grad_x_ptr + offs_m[:, None] * stride_grad_xm + offs_cn[None, :] * stride_grad_xCn + ) + tl.store(grad_x_ptrs, grad_x, mask=mask_m[:, None] & mask_cn[None, :]) + + grad_H_post = tl.reshape(grad_H_post_acc, (BLOCK_SIZE_M * n,)) # (BLOCK_SIZE_M * n) offs_grad_H_post = pid_m * BLOCK_SIZE_M * n + tl.arange(0, BLOCK_SIZE_M * n) grad_H_post_ptrs = grad_H_post_ptr + offs_grad_H_post - tl.atomic_add(grad_H_post_ptrs, grad_H_post, mask=offs_grad_H_post < M * n, sem="relaxed") - - x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_cn[None, :] * stride_xCn - x = tl.load( - x_ptrs, mask=mask_m[:, None] & mask_cn[None, :], other=0.0 - ) # (BLOCK_SIZE_M, BLOCK_SIZE_C*n) - x = tl.reshape(x, (BLOCK_SIZE_M, BLOCK_SIZE_C, n)) # (BLOCK_SIZE_M, BLOCK_SIZE_C, n) - # grad_H_res = x.T @ grad_output: (BLOCK_SIZE_M, n, BLOCK_SIZE_C) @ (BLOCK_SIZE_M, BLOCK_SIZE_C, n) = (BLOCK_SIZE_M, n, n) - grad_H_res = tl.dot( - tl.trans(x, (0, 2, 1)), grad_out, input_precision=precision, out_dtype=tl.float32 - ) # (BLOCK_SIZE_M, n, n) - grad_H_res = tl.reshape(grad_H_res, (BLOCK_SIZE_M * n * n,)) # (BLOCK_SIZE_M * n * n) + grad_H_res = tl.reshape(grad_H_res_acc, (BLOCK_SIZE_M * n * n,)) # (BLOCK_SIZE_M * n * n) offs_grad_H_res = pid_m * BLOCK_SIZE_M * n * n + tl.arange(0, BLOCK_SIZE_M * n * n) grad_H_res_ptrs = grad_H_res_ptr + offs_grad_H_res - tl.atomic_add( - grad_H_res_ptrs, grad_H_res.to(tl.float32), mask=offs_grad_H_res < M * n * n, sem="relaxed" - ) - - grad_out_reshape = tl.reshape( - grad_out, (BLOCK_SIZE_M, BLOCK_SIZE_C, 2, 2) - ) # (BLOCK_SIZE_M, BLOCK_SIZE_C, 2, 2) - grad_out01, grad_out23 = tl.split( - grad_out_reshape - ) # (BLOCK_SIZE_M, BLOCK_SIZE_C, 2), (BLOCK_SIZE_M, BLOCK_SIZE_C, 2) - grad_out0, grad_out1 = tl.split( - grad_out01 - ) # (BLOCK_SIZE_M, BLOCK_SIZE_C), (BLOCK_SIZE_M, BLOCK_SIZE_C) - grad_out2, grad_out3 = tl.split( - grad_out23 - ) # (BLOCK_SIZE_M, BLOCK_SIZE_C), (BLOCK_SIZE_M, BLOCK_SIZE_C) - - # grad_f = grad_output @ H_post.T: (BLOCK_SIZE_M, 1, n) @ (BLOCK_SIZE_M, n, BLOCK_SIZE_C) = (BLOCK_SIZE_M, 1, BLOCK_SIZE_C) - # Triton doesn't support dot prod with inner dimension < 16, so we need to hack this: - # = grad_out[:, :, 0] @ H_post.T[:, 0, :] (BLOCK_SIZE_M, BLOCK_SIZE_C, 1) @ (BLOCK_SIZE_M, 1, 1) - # + grad_out[:, :, 1] @ H_post.T[:, 1, :] - # + grad_out[:, :, 2] @ H_post.T[:, 2, :] - # + grad_out[:, :, 3] @ H_post.T[:, 3, :] - # where H_post.T[:, i, :] = H_post[:, :, i] - H_post = tl.reshape(H_post, (BLOCK_SIZE_M, 2, 2)) - H_post01, H_post23 = tl.split(H_post) # (BLOCK_SIZE_M, 2), (BLOCK_SIZE_M, 2) - H_post0, H_post1 = tl.split(H_post01) # (BLOCK_SIZE_M,), (BLOCK_SIZE_M,) - H_post2, H_post3 = tl.split(H_post23) # (BLOCK_SIZE_M,), (BLOCK_SIZE_M,) - - grad_f_acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_C), dtype=tl.float32) - # (BLOCK_SIZE_M, BLOCK_SIZE_C) * (BLOCK_SIZE_M, 1) -> (BLOCK_SIZE_M, BLOCK_SIZE_C) - grad_f_acc = tl.fma(grad_out0, H_post0[:, None], grad_f_acc) - grad_f_acc = tl.fma(grad_out1, H_post1[:, None], grad_f_acc) - grad_f_acc = tl.fma(grad_out2, H_post2[:, None], grad_f_acc) - grad_f_acc = tl.fma(grad_out3, H_post3[:, None], grad_f_acc) - grad_f = grad_f_acc.to(f.dtype) - - grad_f_ptrs = grad_f_ptr + offs_m[:, None] * stride_grad_fm + offs_c[None, :] * stride_grad_fc - tl.store(grad_f_ptrs, grad_f, mask=mask_m[:, None] & mask_c[None, :]) - - grad_bias = tl.sum(grad_f_acc, axis=0) # (BLOCK_SIZE_C,) - grad_bias_ptrs = grad_bias_ptr + offs_c * stride_grad_bias - tl.atomic_add(grad_bias_ptrs, grad_bias, mask=mask_c, sem="relaxed") - - # grad_x = grad_output @ H_res.T: (BLOCK_SIZE_M, BLOCK_SIZE_C, n) @ (BLOCK_SIZE_M, n, n) = (BLOCK_SIZE_M, n, BLOCK_SIZE_C) - # The inner dim is n=4 which is too small for triton, so we will manually unroll the matmul - # grad_x = grad_out[:, :, 0] @ H_res.T[:, 0, :] - # + grad_out[:, :, 1] @ H_res.T[:, 1, :] - # + grad_out[:, :, 2] @ H_res.T[:, 2, :] - # + grad_out[:, :, 3] @ H_res.T[:, 3, :] - # where H_res.T[:, i, :] = H_res[:, :, i] - # Due to broadcasting, it's equivalent to multiplying each H_res[:, i, :].T with grad_out[:, i, :] - - H_res_reshape = tl.reshape(H_res, (BLOCK_SIZE_M, n, 2, 2)) # (BLOCK_SIZE_M, n, 2, 2) - H_res01, H_res23 = tl.split(H_res_reshape) # (BLOCK_SIZE_M, n, 2), (BLOCK_SIZE_M, n, 2) - H_res0, H_res1 = tl.split(H_res01) # (BLOCK_SIZE_M, n), (BLOCK_SIZE_M, n) - H_res2, H_res3 = tl.split(H_res23) # (BLOCK_SIZE_M, n), (BLOCK_SIZE_M, n) - - grad_x_acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_C, n), dtype=tl.float32) - grad_x_acc = tl.fma(grad_out0[:, :, None], H_res0[:, None, :], grad_x_acc) - grad_x_acc = tl.fma(grad_out1[:, :, None], H_res1[:, None, :], grad_x_acc) - grad_x_acc = tl.fma(grad_out2[:, :, None], H_res2[:, None, :], grad_x_acc) - grad_x_acc = tl.fma(grad_out3[:, :, None], H_res3[:, None, :], grad_x_acc) - - grad_x = grad_x_acc.to(x.dtype) - grad_x = tl.reshape(grad_x, (BLOCK_SIZE_M, BLOCK_SIZE_C * n)) # (BLOCK_SIZE_M, BLOCK_SIZE_C*n) - grad_x_ptrs = grad_x_ptr + offs_m[:, None] * stride_grad_xm + offs_cn[None, :] * stride_grad_xCn - tl.store(grad_x_ptrs, grad_x, mask=mask_m[:, None] & mask_cn[None, :]) + # A single C block covers the reduction, so no atomic_add is needed. + tl.store(grad_H_post_ptrs, grad_H_post.to(H_post.dtype), mask=offs_grad_H_post < M * n) + tl.store(grad_H_res_ptrs, grad_H_res.to(H_res.dtype), mask=offs_grad_H_res < M * n * n) diff --git a/transformer_engine/common/util/cuda_runtime.cpp b/transformer_engine/common/util/cuda_runtime.cpp index 2d29f7e06d..d9a27a2cb7 100644 --- a/transformer_engine/common/util/cuda_runtime.cpp +++ b/transformer_engine/common/util/cuda_runtime.cpp @@ -9,6 +9,7 @@ #include "../util/cuda_runtime.h" #include +#include #include #include @@ -29,6 +30,83 @@ namespace { // String with build-time CUDA include path #include "string_path_cuda_include.h" +// Get the runtime directory of the shared library that contains this code +std::filesystem::path shared_library_directory() { + static const char library_anchor = 0; + Dl_info library_info{}; + if (dladdr(static_cast(&library_anchor), &library_info) == 0 || + library_info.dli_fname == nullptr) { + return {}; + } + + std::filesystem::path library_path = library_info.dli_fname; + if (library_path.is_relative()) { + std::error_code error; + library_path = std::filesystem::absolute(library_path, error); + if (error) { + return {}; + } + } + + return library_path.parent_path(); +} + +std::string runtime_cuda_major_version() { + int runtime_version = 0; + // Header discovery is best-effort, so do not throw if the runtime cannot + // report its version. + if (cudaRuntimeGetVersion(&runtime_version) != cudaSuccess || runtime_version <= 0) { + return {}; + } + + return std::to_string(runtime_version / 1000); +} + +std::filesystem::path python_cuda_directory() { + using Path = std::filesystem::path; + + // Find the Python package root from the installed Transformer Engine package. Do not + // assume that the root is named site-packages or dist-packages since valid installs + // may use an arbitrary target directory. + Path te_package_directory = shared_library_directory(); + while (true) { + if (te_package_directory.filename() == "transformer_engine") { + break; + } + + const Path parent = te_package_directory.parent_path(); + if (parent == te_package_directory) { + // Root directory reached + return {}; + } + + te_package_directory = parent; + } + + const Path nvidia_directory = te_package_directory.parent_path() / "nvidia"; + const auto cuda_major_version = runtime_cuda_major_version(); + if (cuda_major_version.empty()) { + return {}; + } + + std::error_code error; + const Path cuda_directory = nvidia_directory / ("cu" + cuda_major_version); + if (std::filesystem::is_directory(cuda_directory, error)) { + return cuda_directory; + } + + // CUDA 12 Python wheels use the older nvidia/cuda_runtime layout. + if (cuda_major_version == "12") { + error.clear(); + const Path legacy_cuda_directory = nvidia_directory / "cuda_runtime"; + if (std::filesystem::is_directory(legacy_cuda_directory, error)) { + return legacy_cuda_directory; + } + } + + return {}; +} + } // namespace #endif // #ifndef __HIP_PLATFORM_AMD__ @@ -203,6 +281,7 @@ const std::string &include_directory(bool required) { std::vector> search_paths = {{"NVTE_CUDA_INCLUDE_DIR", ""}, {"CUDA_HOME", ""}, {"CUDA_DIR", ""}, + {"", python_cuda_directory()}, {"", string_path_cuda_include}, {"", "/usr/local/cuda"}}; #endif diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index 62668251ef..0a3982daf7 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -127,7 +127,7 @@ constexpr bool is_supported_arch() { NVTE_CUDA_ARCH_MATCHES(ptx::FamilySpecific<100>, ptx::FamilySpecific<110>, \ ptx::FamilySpecific<120>) #define ARCH_HAS_STOCHASTIC_ROUNDING \ - NVTE_CUDA_ARCH_MATCHES(ptx::ArchSpecific<100>, ptx::ArchSpecific<103>) + NVTE_CUDA_ARCH_MATCHES(ptx::ArchSpecific<100>, ptx::ArchSpecific<103>, ptx::ArchSpecific<107>) // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-init __device__ __forceinline__ void mbarrier_init(uint64_t *mbar, const uint32_t count) { @@ -394,6 +394,10 @@ __device__ __forceinline__ bf16 exp2f_rcp(e8m0_t biased_exp) { } __device__ __forceinline__ float exp2f(e8m0_t biased_exp) { + // Handle the special case of NaN. + if (biased_exp == 255) return __int_as_float(0x7fffffff); + // 2^-127 is subnormal, so it cannot be built by shifting into the exponent field. + if (biased_exp == 0) return __int_as_float(0x00400000); return __int_as_float(biased_exp << FP32_MANTISSA_BITS); } @@ -621,6 +625,29 @@ __device__ __forceinline__ void mul_cvt_4x(fp4e2m1x4 &out, const Tx2 &in01, cons out = fp4e2m1x4(make_float4(x0, x1, x2, x3)); } +// Software stochastic rounding onto the e2m1 grid, for architectures without +// cvt.rs. Takes 8 random bits per element, which is what +// cvt.rs.satfinite.e2m1x4.f32 takes from its rbits operand. Follows satfinite +// for the edges: NaN becomes positive MAX_NORM and larger magnitudes clamp to +// MAX_NORM with the sign kept. The result is exactly representable in e2m1, so +// packing it is lossless. +__device__ __forceinline__ float stochastic_round_fp4_e2m1(const float x, const uint32_t rbits8) { + constexpr float max_norm = 6.0f; + const float u = static_cast(rbits8 & 0xFFu) * (1.0f / 256.0f); + const float a = fabsf(x); + // Grid step at |x|: 0.5 below 2, 1 in [2, 4), 2 in [4, 6]. + const float step = (a >= 4.0f) ? 2.0f : ((a >= 2.0f) ? 1.0f : 0.5f); + const float t = fmaf(u, step, a); + // The jitter can carry t across one region boundary, so the rounding step + // derives from t. Flooring in units of that step lands on the e2m1 grid in + // every region, and the saturation also maps a NaN t to max_norm, which the + // sign select keeps positive. + const float step_t = (t >= 4.0f) ? 2.0f : ((t >= 2.0f) ? 1.0f : 0.5f); + const float inv_step_t = (t >= 4.0f) ? 0.5f : ((t >= 2.0f) ? 1.0f : 2.0f); + const float q = fminf(floorf(t * inv_step_t) * step_t, max_norm); + return copysignf(q, (x != x) ? 1.0f : x); +} + __device__ __forceinline__ fp4e2m1x4 mul_cvt_bf16_to_fp4_4x_with_stochastic_rounding( const uint64_t in_4x, const float2 scale, const uint32_t rbits) { uint16_t out_4x = 0; @@ -655,9 +682,14 @@ __device__ __forceinline__ fp4e2m1x4 mul_cvt_bf16_to_fp4_4x_with_stochastic_roun : "=h"(out_4x) : "l"(in_4x), "l"(reinterpret_cast(scale)), "r"(rbits)); } else { - NVTE_DEVICE_ERROR( - "FP4 cvt PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); + // mul.f32x2 above applies scale.x to the even elements and scale.y to the odd ones. + const bf16 *vals = reinterpret_cast(&in_4x); + const float q0 = stochastic_round_fp4_e2m1(static_cast(vals[0]) * scale.x, rbits); + const float q1 = stochastic_round_fp4_e2m1(static_cast(vals[1]) * scale.y, rbits >> 8); + const float q2 = stochastic_round_fp4_e2m1(static_cast(vals[2]) * scale.x, rbits >> 16); + const float q3 = stochastic_round_fp4_e2m1(static_cast(vals[3]) * scale.y, rbits >> 24); + const fp4e2m1x4 packed(make_float4(q0, q1, q2, q3)); + out_4x = *reinterpret_cast(&packed); } #else NVTE_DEVICE_ERROR( @@ -818,9 +850,12 @@ __device__ __forceinline__ fp4e2m1x4 mul_cvt_fp32_to_fp4_4x_with_stochastic_roun "l"(reinterpret_cast(in23)), "l"(reinterpret_cast(scale)), "r"(rbits)); } else { - NVTE_DEVICE_ERROR( - "FP4 cvt PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); + const float q0 = stochastic_round_fp4_e2m1(in01.x * scale.x, rbits); + const float q1 = stochastic_round_fp4_e2m1(in01.y * scale.y, rbits >> 8); + const float q2 = stochastic_round_fp4_e2m1(in23.x * scale.x, rbits >> 16); + const float q3 = stochastic_round_fp4_e2m1(in23.y * scale.y, rbits >> 24); + const fp4e2m1x4 packed(make_float4(q0, q1, q2, q3)); + out_4x = *reinterpret_cast(&packed); } return *reinterpret_cast(&out_4x); #endif @@ -1062,9 +1097,26 @@ __device__ __forceinline__ uint32_t mul_cvt_bf16_to_fp4_8x_stochastic_rounding( NVTE_DEVICE_ERROR("Not supported scaling coefficient type."); } } else { - NVTE_DEVICE_ERROR( - "FP4 cvt PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); + constexpr bool known_coeff = std::is_same::value || + std::is_same::value; + if constexpr (known_coeff) { + const float coeff = static_cast(scaling_coefficient); + const bf16 *vals03 = reinterpret_cast(&in03); + const bf16 *vals47 = reinterpret_cast(&in47); + float q[8]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + q[i] = stochastic_round_fp4_e2m1(static_cast(vals03[i]) * coeff, rbits03 >> (8 * i)); + q[i + 4] = + stochastic_round_fp4_e2m1(static_cast(vals47[i]) * coeff, rbits47 >> (8 * i)); + } + const fp4e2m1x4 lo(make_float4(q[0], q[1], q[2], q[3])); + const fp4e2m1x4 hi(make_float4(q[4], q[5], q[6], q[7])); + out_8x = static_cast(*reinterpret_cast(&lo)) | + (static_cast(*reinterpret_cast(&hi)) << 16); + } else { + NVTE_DEVICE_ERROR("Not supported scaling coefficient type."); + } } return out_8x; } diff --git a/transformer_engine/debug/pytorch/debug_quantization.py b/transformer_engine/debug/pytorch/debug_quantization.py index ed5fdd4660..107dd7a373 100644 --- a/transformer_engine/debug/pytorch/debug_quantization.py +++ b/transformer_engine/debug/pytorch/debug_quantization.py @@ -676,20 +676,31 @@ def size(self, *args): def update_usage(self, rowwise_usage: bool = None, columnwise_usage: bool = None): """Update usage of the tensor.""" - if self.rowwise_gemm_tensor is not self.columnwise_gemm_tensor: - # If the same object is used both for rowwise and columnwise gemms, - # there is no benefit in erasing the usage of one of them. - # And there are scenarios when not deleting the usage of one of them is needed. - # For example when we want to recreate columnwise from rowwise. + if self.rowwise_gemm_tensor is self.columnwise_gemm_tensor: + if isinstance(self.rowwise_gemm_tensor, QuantizedTensor): + self.rowwise_gemm_tensor.update_usage(rowwise_usage, columnwise_usage) + else: + # Each tensor owns only the representation for its GEMM direction. + # Validate requests before dropping either representation so a failed + # mixed drop/enable request does not partially mutate this wrapper. + if rowwise_usage and self.rowwise_gemm_tensor is None: + raise RuntimeError( + "Cannot recreate rowwise tensor from columnwise tensor in debug mode." + ) + if columnwise_usage and self.columnwise_gemm_tensor is None: + raise RuntimeError( + "Cannot recreate columnwise tensor from rowwise tensor in debug mode." + ) + if rowwise_usage is False: self.rowwise_gemm_tensor = None + elif rowwise_usage and isinstance(self.rowwise_gemm_tensor, QuantizedTensor): + self.rowwise_gemm_tensor.update_usage(rowwise_usage=True) + if columnwise_usage is False: self.columnwise_gemm_tensor = None - - if isinstance(self.rowwise_gemm_tensor, QuantizedTensor): - self.rowwise_gemm_tensor.update_usage(rowwise_usage, columnwise_usage) - if isinstance(self.columnwise_gemm_tensor, QuantizedTensor): - self.columnwise_gemm_tensor.update_usage(rowwise_usage, columnwise_usage) + elif columnwise_usage and isinstance(self.columnwise_gemm_tensor, QuantizedTensor): + self.columnwise_gemm_tensor.update_usage(columnwise_usage=True) if rowwise_usage and self.rowwise_gemm_tensor is None: raise RuntimeError( @@ -698,7 +709,7 @@ def update_usage(self, rowwise_usage: bool = None, columnwise_usage: bool = None if columnwise_usage and self.columnwise_gemm_tensor is None: raise RuntimeError( - "Cannot recreate columnwise tensor from rowwise tensor is debug mode." + "Cannot recreate columnwise tensor from rowwise tensor in debug mode." ) @property diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index 1dde00814d..0afa6faafb 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -342,6 +342,7 @@ def is_fused_attn_kernel_available( head_dim_qk, head_dim_v, window_size: Optional[Tuple[int, int]] = None, + return_max_logit: bool = False, ): """ To check whether the fused attention kernel is supported @@ -365,6 +366,7 @@ def make_helper(attn_mask_type): head_dim_qk, head_dim_v, window_size_tuple, + return_max_logit, ) return make_helper(attn_mask_type).is_fused_attn_kernel_available() @@ -1056,6 +1058,7 @@ def _legacy_fused_attn( context_parallel_causal_load_balanced: bool = False, context_parallel_axis: str = "", softmax_offset: Optional[jnp.ndarray] = None, + return_max_logit: bool = False, ): """ Perform non-THD (non-packed) cuDNN fused attention. @@ -1087,8 +1090,15 @@ def _legacy_fused_attn( context_parallel_causal_load_balanced (bool): Indicates the sequences are ordered for causal mask load balancing when running context parallelism. context_parallel_axis (str): The name of the context parallel axis. + softmax_offset (Optional[jnp.ndarray]): An optional learnable softmax offset tensor with shape + [1, num_heads, 1, 1]. Used when softmax_type is AttnSoftmaxType.LEARNABLE_SOFTMAX. + return_max_logit (bool): If True, also return per-head maximum attention logits + with shape ``[h]``. Returns: - (jnp.ndarray): The output tensor from the fused attention. + jnp.ndarray: + Attention output when ``return_max_logit`` is False. + tuple[jnp.ndarray, jnp.ndarray]: + ``(output, max_logit)`` when ``return_max_logit`` is True. """ assert ( not qkv_layout.is_thd() @@ -1142,6 +1152,7 @@ def _legacy_fused_attn( context_parallel_strategy=context_parallel_strategy, context_parallel_causal_load_balanced=context_parallel_causal_load_balanced, context_parallel_axis=context_parallel_axis, + return_max_logit=return_max_logit, ) return output @@ -1167,6 +1178,7 @@ def fused_attn_thd( context_parallel_causal_load_balanced: bool = False, context_parallel_axis: str = "", softmax_offset: Optional[jnp.ndarray] = None, + return_max_logit: bool = False, ): """ Deprecated THD fused attn, please use fusd_attn with SequenceDescriptor @@ -1221,12 +1233,16 @@ def fused_attn_thd( context_parallel_strategy=context_parallel_strategy, context_parallel_causal_load_balanced=context_parallel_causal_load_balanced, context_parallel_axis=context_parallel_axis, + return_max_logit=return_max_logit, ) return output -@partial(jax.custom_vjp, nondiff_argnums=(5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18)) +@partial( + jax.custom_vjp, + nondiff_argnums=(5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19), +) def _fused_attn( qkv: Tuple[jnp.ndarray, ...], bias: Optional[jnp.ndarray], @@ -1247,6 +1263,7 @@ def _fused_attn( context_parallel_axis: str, context_checkpoint_name: str = "context", stripe_size: int | None = None, + return_max_logit: bool = False, ): output, _ = _fused_attn_fwd_rule( qkv, @@ -1268,6 +1285,7 @@ def _fused_attn( context_parallel_axis, context_checkpoint_name=context_checkpoint_name, stripe_size=stripe_size, + return_max_logit=return_max_logit, ) return output @@ -1292,8 +1310,9 @@ def _fused_attn_fwd_rule( context_parallel_axis, context_checkpoint_name, stripe_size, + return_max_logit, ): - output, softmax_aux, rng_state = tex.fused_attn_fwd( + output, softmax_aux, rng_state, max_logit = tex.fused_attn_fwd( qkv, bias, softmax_offset, @@ -1312,11 +1331,14 @@ def _fused_attn_fwd_rule( context_parallel_causal_load_balanced=context_parallel_causal_load_balanced, context_parallel_axis=context_parallel_axis, stripe_size=stripe_size, + return_max_logit=return_max_logit, ) output = checkpoint_name(output, context_checkpoint_name) softmax_aux = checkpoint_name(softmax_aux, context_checkpoint_name) rng_state = checkpoint_name(rng_state, context_checkpoint_name) - return output, ( + max_logit = checkpoint_name(max_logit, context_checkpoint_name) + attn_output = (output, max_logit) if return_max_logit else output + return attn_output, ( qkv, bias, sequence_descriptor, @@ -1342,10 +1364,13 @@ def _fused_attn_bwd_rule( context_parallel_axis, context_checkpoint_name, stripe_size, + return_max_logit, ctx, dz, ): del context_checkpoint_name + if return_max_logit: + dz, _ = dz ( qkv, bias, @@ -1471,6 +1496,7 @@ def fused_attn( score_mod_bprop: Optional[Callable] = None, score_mod_tensors: Optional[Mapping[str, Any]] = None, score_mod_bprop_tensors: Optional[Mapping[str, Any]] = None, + return_max_logit: bool = False, ): """ Perform cuDNN fused attention. @@ -1527,8 +1553,13 @@ def fused_attn( non-differentiable auxiliary inputs. score_mod_bprop_tensors (Optional[Mapping[str, Any]]): Additional tensors or Python/NumPy scalars made available to `score_mod_bprop`. + return_max_logit (bool): If True, also return per-head maximum attention logits + with shape ``[h]``. Returns: - (jnp.ndarray): The output tensor from the fused attention. + jnp.ndarray: + Attention output when ``return_max_logit`` is False. + tuple[jnp.ndarray, jnp.ndarray]: + ``(output, max_logit)`` when ``return_max_logit`` is True. Examples (non-THD, also known as non-packed): >>> # q_segment_ids = [[1, 1, 1, 0], [1, 1, 0, 0]], 0 means padded tokens @@ -1578,6 +1609,8 @@ def fused_attn( # escaping from a graph-building helper mid-trace. if is_hip_extension(): raise NotImplementedError("score_mod fused attention is not supported on ROCm.") + if return_max_logit: + raise ValueError("return_max_logit is not supported with score_mod fused_attn.") tex.validate_fused_attn_score_mod( qkv, bias, @@ -1637,6 +1670,7 @@ def fused_attn( context_parallel_causal_load_balanced=context_parallel_causal_load_balanced, context_parallel_axis=context_parallel_axis, softmax_offset=softmax_offset, + return_max_logit=return_max_logit, ) if max_segments_per_seq > 1 and not qkv_layout.is_thd(): warnings.warn( @@ -1667,5 +1701,6 @@ def fused_attn( context_parallel_axis=context_parallel_axis, context_checkpoint_name=context_checkpoint_name, stripe_size=stripe_size, + return_max_logit=return_max_logit, ) return output diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index d2bab63751..84a3c87247 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -78,6 +78,7 @@ "cp_axis", "cp_striped_window_size", "stripe_size", + "return_max_logit", ], ) @dataclass(frozen=True) @@ -102,6 +103,7 @@ class _FusedAttnConfig: stripe_size: ( int | None ) # Only for CP + Striped. For Ring P2P, stripe_size=1 only.For AG, stripe_size>=1. + return_max_logit: bool = False @dataclass(frozen=True) @@ -125,6 +127,7 @@ class FusedAttnHelper: head_dim_qk: int head_dim_v: int window_size: Tuple[int, int] + return_max_logit: bool = False def is_fused_attn_kernel_available(self): """Check if there is available fused attention kernel""" @@ -149,6 +152,7 @@ def get_fused_attn_backend(self): self.head_dim_v, self.window_size[0], self.window_size[1], + self.return_max_logit, not self.is_non_deterministic_allowed(), ) @@ -360,6 +364,7 @@ def abstract( q_head_dim, v_head_dim, config.window_size, + config.return_max_logit, ).get_fused_attn_backend() if not is_hip_extension(): @@ -393,6 +398,17 @@ def abstract( else: raise ValueError(f"Unsupported {backend=}") softmax_aux_aval = q_aval.update(shape=softmax_shape, dtype=softmax_dtype) + if config.return_max_logit: + # cuDNN Max is row-wise over S_kv. Dense and SM120 THD use + # [..., H, S_q, 1]; cuDNN >= 9.6 non-SM120 THD uses [..., S_q, H, 1]. + # Both raw layouts are reduced to the public per-head [H] result below. + if FusedAttnFwdPrimitive._uses_thd_ragged_max_tensor(config): + max_tensor_shape = (*batch_shape, q_max_seqlen, attn_heads, 1) + else: + max_tensor_shape = (*batch_shape, attn_heads, q_max_seqlen, 1) + else: + max_tensor_shape = (0,) + max_tensor_aval = q_aval.update(shape=max_tensor_shape, dtype=softmax_dtype) # JAX does not enable 64-bit int by default so we get XLA to allocate x8 memory with # 32-bit unsigned int to get the buffer size we need in the C++ kernel @@ -439,6 +455,7 @@ def abstract( config.max_segments_per_seq, config.window_size[0], config.window_size[1], + config.return_max_logit, bottom_right_diagonal, ) wkspace_aval = q_aval.update( @@ -459,17 +476,19 @@ def abstract( f" {softmax_offset_aval.shape}" ) - return out_aval, softmax_aux_aval, rng_state_aval, wkspace_aval + return out_aval, softmax_aux_aval, max_tensor_aval, rng_state_aval, wkspace_aval @staticmethod def outer_abstract(*args, **kwargs): """ Fused attention fwd outer primitive abstract """ - out_aval, softmax_aux_aval, rng_state_aval, _ = FusedAttnFwdPrimitive.abstract( + out_aval, softmax_aux_aval, _, rng_state_aval, _ = FusedAttnFwdPrimitive.abstract( *args, **kwargs ) - return out_aval, softmax_aux_aval, rng_state_aval + max_logit_shape = (out_aval.shape[-2],) if kwargs["config"].return_max_logit else (0,) + max_logit_aval = out_aval.update(shape=max_logit_shape, dtype=out_aval.dtype) + return out_aval, softmax_aux_aval, rng_state_aval, max_logit_aval @staticmethod def lowering( @@ -553,6 +572,7 @@ def lowering( mask_type=int(config.attn_mask_type.value), qkv_layout=int(config.qkv_layout.value), is_training=config.is_training, + return_max_logit=config.return_max_logit, deterministic=not FusedAttnHelper.is_non_deterministic_allowed(), window_size_left=window_size_left, window_size_right=window_size_right, @@ -596,6 +616,9 @@ def impl( config.max_segments_per_seq, ) ) + raw_q_seqlen = q_seqlen + raw_q_seq_offsets = q_seq_offsets + if config.qkv_layout.is_thd(): def _fix_len_take(x, condition, fill_value=-1): @@ -652,7 +675,7 @@ def convert_to_2d(offsets, batch, max_seqlen): q_cu_seqlen = generate_cu_seqlen(q_seqlen.flatten()) kv_cu_seqlen = generate_cu_seqlen(kv_seqlen.flatten()) - output, softmax_aux, rng_state, _ = FusedAttnFwdPrimitive.inner_primitive.bind( + output, softmax_aux, max_tensor, rng_state, _ = FusedAttnFwdPrimitive.inner_primitive.bind( q, k, v, @@ -669,7 +692,91 @@ def convert_to_2d(offsets, batch, max_seqlen): _kv_segment_pos, config=config, ) - return output, softmax_aux, rng_state + # Reduce cuDNN's raw Max tensor to TE's public per-head [H] max_logit. + max_logit = FusedAttnFwdPrimitive._reduce_max_logit( + max_tensor, output, raw_q_seqlen, raw_q_seq_offsets, config + ) + return output, softmax_aux, rng_state, max_logit + + @staticmethod + def _reduce_max_logit(max_tensor, output, q_seqlen, q_seq_offsets, config): + """Reduce cuDNN's row-wise Max tensor to the public per-head max_logit. + + Dense and SM120 THD use ``[..., H, S_q, 1]``; cuDNN >= 9.6 non-SM120 + THD uses ``[..., S_q, H, 1]``. A rank-3 THD result is ``[T_q, H, 1]``. + All layouts reduce to ``[H]``. Static THD buffers can contain invalid query + rows, so those rows are masked before reduction. + """ + if not config.return_max_logit: + return jnp.zeros((0,), dtype=output.dtype) + + uses_thd_ragged_max_tensor = FusedAttnFwdPrimitive._uses_thd_ragged_max_tensor(config) + if config.qkv_layout.is_thd() and max_tensor.ndim == 4: + # Dense BSHD Max rows are expected to be masked by cuDNN before TE reduces them. + # THD Max can include static holes/unwritten rows, so mask valid query rows here. + q_seqlen = jnp.where(q_seqlen > 0, q_seqlen, 0) + q_seq_offsets = jnp.where(q_seq_offsets >= 0, q_seq_offsets, -1) + num_segments = min(q_seqlen.shape[-1], q_seq_offsets.shape[-1]) + q_seqlen = q_seqlen[..., :num_segments] + q_seq_offsets = q_seq_offsets[..., :num_segments] + token_idx = jnp.arange(output.shape[-3], dtype=q_seq_offsets.dtype) + valid = jnp.any( + (q_seq_offsets[..., None] >= 0) + & (token_idx >= q_seq_offsets[..., None]) + & (token_idx < (q_seq_offsets[..., None] + q_seqlen[..., None])), + axis=-2, + ) + if uses_thd_ragged_max_tensor: + max_tensor = jnp.where(valid[:, :, None, None], max_tensor, -jnp.inf) + else: + max_tensor = jnp.where(valid[:, None, :, None], max_tensor, -jnp.inf) + + if max_tensor.ndim == 3: + amax_dims = (0, 2) + elif uses_thd_ragged_max_tensor: + amax_dims = (0, 1, 3) + else: + amax_dims = (0, 2, 3) + return jnp.max(max_tensor, axis=amax_dims).astype(output.dtype) + + @staticmethod + def _uses_thd_ragged_max_tensor(config): + """Return whether cuDNN writes THD Max with BSH-like ragged-stats layout.""" + return ( + config.qkv_layout.is_thd() + and get_cudnn_version() >= (9, 6, 0) + and 120 not in get_all_device_compute_capability() + ) + + @staticmethod + def _empty_or_neg_inf_max_logit(head, dtype, config): + """Return the neutral value for per-head max_logit accumulation.""" + if config.return_max_logit: + return jnp.full((head,), -jnp.inf, dtype=dtype) + return jnp.zeros((0,), dtype=dtype) + + @staticmethod + def _max_logit_reduce_axes(mesh, max_logit_sharding): + """Return mesh axes to reduce while preserving max_logit's head sharding.""" + # max_logit is [H], so axes that shard H (typically TP) are preserved. + # Axes for collapsed dimensions such as batch/sequence (DP/CP) must pmax. + head_axes = set() + for axis in max_logit_sharding.spec: + if axis is None: + continue + if isinstance(axis, tuple): + head_axes.update(axis) + else: + head_axes.add(axis) + return tuple(axis for axis in mesh.axis_names if axis not in head_axes) + + @staticmethod + def _reduce_max_logit_across_mesh(max_logit, mesh, reduce_axes, config): + """Reduce max_logit across mesh axes absent from the [H] result.""" + if config.return_max_logit: + for axis in reduce_axes: + max_logit = lax_paral_op(max_logit, lax.pmax, axis, mesh=mesh) + return max_logit @staticmethod def batcher(batched_args, batch_dims, *, config): @@ -681,7 +788,8 @@ def batcher(batched_args, batch_dims, *, config): q_bdim, _, _, _, _, seed_bdim, *_ = batch_dims # Pass through; segment_ids/segment_pos may have different batch dims (e.g. vmapped ids, # replicated pos). get_seqlens_and_offsets() in attention.py handles conversion without expanding. - out_bdims = q_bdim, q_bdim, seed_bdim + max_logit_bdim = q_bdim if config.return_max_logit else None + out_bdims = q_bdim, q_bdim, seed_bdim, max_logit_bdim return ( FusedAttnFwdPrimitive.outer_primitive.bind(*batched_args, config=config), out_bdims, @@ -735,12 +843,16 @@ def infer_sharding_from_operands(config, mesh, arg_infos, result_infos): raise ValueError(f"Unsupported {config.qkv_layout=}") rng_state_sharding = NamedSharding(mesh, PartitionSpec(get_all_mesh_axes(), None)) - return (out_sharding, softmax_aux_sharding, rng_state_sharding) + max_logit_sharding = NamedSharding( + mesh, PartitionSpec(q_spec[-2] if config.return_max_logit else None) + ) + return (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding) @staticmethod def partition(config, mesh, arg_infos, result_infos): out_sharding = result_infos[0].sharding softmax_aux_sharding = result_infos[1].sharding + max_logit_sharding = result_infos[3].sharding rng_state_sharding = seed_sharding = NamedSharding( mesh, PartitionSpec(get_all_mesh_axes(), None) ) @@ -749,8 +861,23 @@ def partition(config, mesh, arg_infos, result_infos): arg_shardings[-1] = arg_shardings[-3] arg_shardings[-2] = arg_shardings[-4] arg_shardings = tuple(arg_shardings) - out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding) - impl = partial(FusedAttnFwdPrimitive.impl, config=config) + out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding) + max_logit_reduce_axes = ( + FusedAttnFwdPrimitive._max_logit_reduce_axes(mesh, max_logit_sharding) + if config.return_max_logit + else () + ) + + def impl(*args): + output, softmax_aux, rng_state, max_logit = FusedAttnFwdPrimitive.impl( + *args, config=config + ) + # Globalize the rank-local [H] max across DP/CP while preserving TP head sharding. + max_logit = FusedAttnFwdPrimitive._reduce_max_logit_across_mesh( + max_logit, mesh, max_logit_reduce_axes, config + ) + return output, softmax_aux, rng_state, max_logit + return mesh, impl, out_shardings, arg_shardings @staticmethod @@ -778,8 +905,10 @@ def shardy_sharding_rule(config, mesh, value_types, result_types): else: softmax_aux_sharding = ("…0", "head", "seqlen", "i") + max_logit_sharding = ("head",) if config.return_max_logit else ("max_logit",) return SdyShardingRule( - tuple(input_spec), (out_sharding, softmax_aux_sharding, rng_sharding) + tuple(input_spec), + (out_sharding, softmax_aux_sharding, rng_sharding, max_logit_sharding), ) @@ -1448,6 +1577,7 @@ def get_step_config(self) -> _FusedAttnConfig: cp_axis=self.config.cp_axis, cp_striped_window_size=None, stripe_size=self.config.stripe_size, + return_max_logit=self.config.return_max_logit, ) def get_step_config_for_striped(self, max_seqlen, cp_size) -> _FusedAttnConfig: @@ -1468,6 +1598,7 @@ def get_step_config_for_striped(self, max_seqlen, cp_size) -> _FusedAttnConfig: cp_axis=self.config.cp_axis, cp_striped_window_size=None, stripe_size=self.config.stripe_size, + return_max_logit=self.config.return_max_logit, ) def all_gather_kv(self, k, v): @@ -1838,13 +1969,19 @@ def partition(config, mesh, arg_infos, result_infos): out_sharding = result_infos[0].sharding softmax_aux_sharding = result_infos[1].sharding + max_logit_sharding = result_infos[3].sharding rng_state_sharding = seed_sharding = NamedSharding( mesh, PartitionSpec(get_all_mesh_axes(), None) ) arg_shardings = [arg_i.sharding for arg_i in arg_infos] arg_shardings[5] = seed_sharding arg_shardings = tuple(arg_shardings) - out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding) + out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding) + max_logit_reduce_axes = ( + FusedAttnFwdPrimitive._max_logit_reduce_axes(mesh, max_logit_sharding) + if config.return_max_logit + else () + ) def impl( q, @@ -1892,7 +2029,8 @@ def _cross_attn(idx, q, k, v, bias, softmax_offset, q_seqlen, kv_seqlen, seed): q_seqlen_for_step = q_seqlen / (cp_size * 2) num_kv_chunks = kv_max_seqlen // kv_seqlens_for_rank[sub_idx] kv_seqlen_for_step = (kv_seqlen / (cp_size * 2)) * num_kv_chunks - output, softmax_aux, rng_state = FusedAttnFwdPrimitive.impl( + # max_logit returned here is already reduced to shape [H] + output, softmax_aux, rng_state, max_logit = FusedAttnFwdPrimitive.impl( q_split[sub_idx], k_unmasked, v_unmasked, @@ -1909,13 +2047,15 @@ def _cross_attn(idx, q, k, v, bias, softmax_offset, q_seqlen, kv_seqlen, seed): _kv_segment_pos, config=helper.get_step_config(), ) - results.append((output, softmax_aux, rng_state)) + results.append((output, softmax_aux, rng_state, max_logit)) output = jnp.concatenate((results[0][0], results[1][0]), axis=1) softmax_aux = jnp.concatenate((results[0][1], results[1][1]), axis=2) rng_state = results[1][2] # Use the final RNG state + # Rank-local [H] max across both local dual-chunk query pieces. + max_logit = jnp.maximum(results[0][3], results[1][3]) - return output, softmax_aux, rng_state + return output, softmax_aux, rng_state, max_logit k_ag, v_ag = helper.all_gather_kv(k, v) @@ -1926,7 +2066,12 @@ def _cross_attn(idx, q, k, v, bias, softmax_offset, q_seqlen, kv_seqlen, seed): for idx in range(cp_size) ] - return lax.switch(cp_rank, functions) + output, softmax_aux, rng_state, max_logit = lax.switch(cp_rank, functions) + # Globalize the rank-local [H] max across DP/CP while preserving TP head sharding. + max_logit = FusedAttnFwdPrimitive._reduce_max_logit_across_mesh( + max_logit, mesh, max_logit_reduce_axes, config + ) + return output, softmax_aux, rng_state, max_logit return mesh, impl, out_shardings, arg_shardings @@ -2131,13 +2276,19 @@ def partition(config, mesh, arg_infos, result_infos): out_sharding = result_infos[0].sharding softmax_aux_sharding = result_infos[1].sharding + max_logit_sharding = result_infos[3].sharding rng_state_sharding = seed_sharding = NamedSharding( mesh, PartitionSpec(get_all_mesh_axes(), None) ) arg_shardings = [arg_i.sharding for arg_i in arg_infos] arg_shardings[5] = seed_sharding arg_shardings = tuple(arg_shardings) - out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding) + out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding) + max_logit_reduce_axes = ( + FusedAttnFwdPrimitive._max_logit_reduce_axes(mesh, max_logit_sharding) + if config.return_max_logit + else () + ) def impl( q, @@ -2201,7 +2352,7 @@ def _cross_attn( max_segments_per_seq=adjusted_max_segments_per_seq, ) - output, softmax_aux, rng_state = FusedAttnFwdPrimitive.impl( + output, softmax_aux, rng_state, max_logit = FusedAttnFwdPrimitive.impl( q, # sharded for rank k, # ag v, # ag @@ -2220,7 +2371,7 @@ def _cross_attn( max_seqlen=kv_max_seqlen, cp_size=cp_size ), ) - return output, softmax_aux, rng_state + return output, softmax_aux, rng_state, max_logit # AG the k, v, kv_segment_ids and kv_segment_pos k_ag, v_ag = helper.all_gather_kv(k, v) @@ -2241,7 +2392,12 @@ def _cross_attn( ) for _ in range(cp_size) ] - return lax.switch(cp_rank, functions) + output, softmax_aux, rng_state, max_logit = lax.switch(cp_rank, functions) + # Globalize the rank-local [H] max across DP/CP while preserving TP head sharding. + max_logit = FusedAttnFwdPrimitive._reduce_max_logit_across_mesh( + max_logit, mesh, max_logit_reduce_axes, config + ) + return output, softmax_aux, rng_state, max_logit return mesh, impl, out_shardings, arg_shardings @@ -2514,6 +2670,7 @@ def get_step_config(self, attn_mask_type) -> _FusedAttnConfig: cp_axis=self.config.cp_axis, cp_striped_window_size=None, stripe_size=self.config.stripe_size, + return_max_logit=self.config.return_max_logit, ) def stack_kv(self, k, v): @@ -2581,6 +2738,7 @@ def partition(config, mesh, arg_infos, result_infos): out_sharding = result_infos[0].sharding softmax_aux_sharding = result_infos[1].sharding + max_logit_sharding = result_infos[3].sharding rng_state_sharding = seed_sharding = NamedSharding( mesh, PartitionSpec(get_all_mesh_axes(), None) ) @@ -2590,7 +2748,12 @@ def partition(config, mesh, arg_infos, result_infos): arg_shardings[-1] = arg_shardings[-3] arg_shardings[-2] = arg_shardings[-4] arg_shardings = tuple(arg_shardings) - out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding) + out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding) + max_logit_reduce_axes = ( + FusedAttnFwdPrimitive._max_logit_reduce_axes(mesh, max_logit_sharding) + if config.return_max_logit + else () + ) def ring_attn_fwd_impl( q, @@ -2628,9 +2791,10 @@ def ring_attn_fwd_impl( # support dropout currently. rng_state_shape = (seed.shape[0], *result_infos[2].shape[1:]) rng_state = jnp.zeros(rng_state_shape).astype(result_infos[2].dtype) + max_logit = FusedAttnFwdPrimitive._empty_or_neg_inf_max_logit(head, q.dtype, config) def scan_kv_block(idx, carry): - kv, output, softmax_aux = carry + kv, output, softmax_aux, max_logit = carry # Send KV block to next step so we can overlap compute. kv_next = helper.permute_kv(kv, cp_perm) @@ -2638,24 +2802,26 @@ def scan_kv_block(idx, carry): def mask_compute(attn_mask_type): q_seqlen_per_step = helper.adjust_seqlen(q_seqlen, q_max_seqlen, idx) kv_seqlen_per_step = helper.adjust_seqlen(kv_seqlen, kv_max_seqlen, idx) - output_per_step, softmax_aux_per_step, _ = FusedAttnFwdPrimitive.impl( - q, - kv, - _not_used, - bias, - _softmax_offset, - seed, - q_seqlen_per_step, - kv_seqlen_per_step, - q_seq_offsets, - k_seq_offsets, - _q_segment_ids, - _kv_segment_ids, - _q_segment_pos, - _kv_segment_pos, - config=helper.get_step_config(attn_mask_type), + output_per_step, softmax_aux_per_step, _, max_logit_per_step = ( + FusedAttnFwdPrimitive.impl( + q, + kv, + _not_used, + bias, + _softmax_offset, + seed, + q_seqlen_per_step, + kv_seqlen_per_step, + q_seq_offsets, + k_seq_offsets, + _q_segment_ids, + _kv_segment_ids, + _q_segment_pos, + _kv_segment_pos, + config=helper.get_step_config(attn_mask_type), + ) ) - return output_per_step, softmax_aux_per_step + return output_per_step, softmax_aux_per_step, max_logit_per_step causal_mask_compute = partial(mask_compute, AttnMaskType.CAUSAL_MASK) no_mask_compute = partial(mask_compute, AttnMaskType.NO_MASK) @@ -2664,45 +2830,49 @@ def half_kv_no_mask_compute(): q_seqlen_per_step = helper.adjust_seqlen(q_seqlen, q_max_seqlen, idx) kv_seqlen_per_step = helper.adjust_seqlen(kv_seqlen, kv_max_seqlen, idx) // 2 kv_part = lax.slice_in_dim(kv, 0, kv.shape[1] // 2, axis=1) - output_per_step, softmax_aux_per_step, _ = FusedAttnFwdPrimitive.impl( - q, - kv_part, - _not_used, - bias, - _softmax_offset, - seed, - q_seqlen_per_step, - kv_seqlen_per_step, - q_seq_offsets, - k_seq_offsets, - _q_segment_ids, - _kv_segment_ids, - _q_segment_pos, - _kv_segment_pos, - config=helper.get_step_config(AttnMaskType.NO_MASK), + output_per_step, softmax_aux_per_step, _, max_logit_per_step = ( + FusedAttnFwdPrimitive.impl( + q, + kv_part, + _not_used, + bias, + _softmax_offset, + seed, + q_seqlen_per_step, + kv_seqlen_per_step, + q_seq_offsets, + k_seq_offsets, + _q_segment_ids, + _kv_segment_ids, + _q_segment_pos, + _kv_segment_pos, + config=helper.get_step_config(AttnMaskType.NO_MASK), + ) ) - return output_per_step, softmax_aux_per_step + return output_per_step, softmax_aux_per_step, max_logit_per_step def half_q_no_mask_compute(): q_seqlen_per_step = helper.adjust_seqlen(q_seqlen, q_max_seqlen, idx) // 2 kv_seqlen_per_step = helper.adjust_seqlen(kv_seqlen, kv_max_seqlen, idx) q_part = lax.slice_in_dim(q, q_max_seqlen // 2, q_max_seqlen, axis=1) - output_per_step, softmax_aux_per_step, _ = FusedAttnFwdPrimitive.impl( - q_part, - kv, - _not_used, - bias, - _softmax_offset, - seed, - q_seqlen_per_step, - kv_seqlen_per_step, - q_seq_offsets, - k_seq_offsets, - _q_segment_ids, - _kv_segment_ids, - _q_segment_pos, - _kv_segment_pos, - config=helper.get_step_config(AttnMaskType.NO_MASK), + output_per_step, softmax_aux_per_step, _, max_logit_per_step = ( + FusedAttnFwdPrimitive.impl( + q_part, + kv, + _not_used, + bias, + _softmax_offset, + seed, + q_seqlen_per_step, + kv_seqlen_per_step, + q_seq_offsets, + k_seq_offsets, + _q_segment_ids, + _kv_segment_ids, + _q_segment_pos, + _kv_segment_pos, + config=helper.get_step_config(AttnMaskType.NO_MASK), + ) ) output_per_step = jnp.concat([jnp.zeros_like(q_part), output_per_step], axis=1) softmax_aux_per_step = jnp.concat( @@ -2712,14 +2882,17 @@ def half_q_no_mask_compute(): ], axis=2, ) - return output_per_step, softmax_aux_per_step + return output_per_step, softmax_aux_per_step, max_logit_per_step def skip_compute(): output_per_step = jnp.zeros_like(q) softmax_aux_per_step = jnp.full( (batch, head, q.shape[1], 1), -jnp.inf, dtype=jnp.float32 ) - return output_per_step, softmax_aux_per_step + max_logit_per_step = FusedAttnFwdPrimitive._empty_or_neg_inf_max_logit( + head, q.dtype, config + ) + return output_per_step, softmax_aux_per_step, max_logit_per_step if config.attn_mask_type == AttnMaskType.CAUSAL_MASK: # This is for nested jax.lax.cond @@ -2730,11 +2903,11 @@ def jax_cond_wrap(): ) return lax.cond((idx <= cp_rank), no_mask_compute, skip_compute) - output_per_step, softmax_aux_per_step = lax.cond( + output_per_step, softmax_aux_per_step, max_logit_per_step = lax.cond( idx == 0, causal_mask_compute, jax_cond_wrap ) else: - output_per_step, softmax_aux_per_step = no_mask_compute() + output_per_step, softmax_aux_per_step, max_logit_per_step = no_mask_compute() def skip_correction(output, softmax_aux, output_per_step, softmax_aux_per_step): # No correction done here but we cast outputs to float32 and perform reduction @@ -2757,19 +2930,25 @@ def correction(output, softmax_aux, output_per_step, softmax_aux_per_step): output_per_step, softmax_aux_per_step, ) + # Running per-head max over all ring steps for this rank. + max_logit = jnp.maximum(max_logit, max_logit_per_step) - return (kv_next, output, softmax_aux) + return (kv_next, output, softmax_aux, max_logit) - carry = (kv, output, softmax_aux) + carry = (kv, output, softmax_aux, max_logit) if helper.use_scanloop(): carry = lax.fori_loop(0, cp_size, scan_kv_block, carry) else: for i in range(0, cp_size): carry = scan_kv_block(i, carry) - (kv, output, softmax_aux) = carry + (kv, output, softmax_aux, max_logit) = carry output = output.astype(q.dtype) - return output, softmax_aux, rng_state + # Globalize the rank-local running [H] max across DP/CP. + max_logit = FusedAttnFwdPrimitive._reduce_max_logit_across_mesh( + max_logit, mesh, max_logit_reduce_axes, config + ) + return output, softmax_aux, rng_state, max_logit return mesh, ring_attn_fwd_impl, out_shardings, arg_shardings @@ -3087,6 +3266,7 @@ def partition(config, mesh, arg_infos, result_infos): out_sharding = result_infos[0].sharding softmax_aux_sharding = result_infos[1].sharding + max_logit_sharding = result_infos[3].sharding rng_state_sharding = seed_sharding = NamedSharding( mesh, PartitionSpec(get_all_mesh_axes(), None) ) @@ -3096,7 +3276,12 @@ def partition(config, mesh, arg_infos, result_infos): arg_shardings[-1] = arg_shardings[-3] arg_shardings[-2] = arg_shardings[-4] arg_shardings = tuple(arg_shardings) - out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding) + out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding) + max_logit_reduce_axes = ( + FusedAttnFwdPrimitive._max_logit_reduce_axes(mesh, max_logit_sharding) + if config.return_max_logit + else () + ) def fwd_impl( q, @@ -3139,9 +3324,10 @@ def fwd_impl( # support dropout currently. rng_state_shape = (seed.shape[0], *result_infos[2].shape[1:]) rng_state = jnp.zeros(rng_state_shape).astype(result_infos[2].dtype) + max_logit = FusedAttnFwdPrimitive._empty_or_neg_inf_max_logit(head, q.dtype, config) def scan_kv_block(idx, carry): - kv, kv_segment_ids, kv_segment_pos, output, softmax_aux = carry + kv, kv_segment_ids, kv_segment_pos, output, softmax_aux, max_logit = carry # TODO(rewang): To check whether we need special handle for the last idx # Send KV block to next step so we can overlap compute. @@ -3179,7 +3365,9 @@ def compute(config): ) else: current_config = subblock_config - output_per_step, softmax_aux_per_step, _ = compute(current_config) + output_per_step, softmax_aux_per_step, _, max_logit_per_step = compute( + current_config + ) softmax_aux_per_step = softmax_aux_per_step.reshape((batch, q_max_seqlen, head, 1)) @@ -3214,18 +3402,32 @@ def correction(output, softmax_aux, output_per_step, softmax_aux_per_step): output_per_step, softmax_aux_per_step, ) + # Running per-head max over all ring steps for this rank. + max_logit = jnp.maximum(max_logit, max_logit_per_step) - return (kv_next, kv_segment_ids_next, kv_segment_pos_next, output, softmax_aux) + return ( + kv_next, + kv_segment_ids_next, + kv_segment_pos_next, + output, + softmax_aux, + max_logit, + ) - carry = (kv, kv_segment_ids, kv_segment_pos, output, softmax_aux) + carry = (kv, kv_segment_ids, kv_segment_pos, output, softmax_aux, max_logit) if helper.use_scanloop(): carry = lax.fori_loop(0, cp_size, scan_kv_block, carry) else: for i in range(0, cp_size): carry = scan_kv_block(i, carry) - (_, _, _, output, softmax_aux) = carry + (_, _, _, output, softmax_aux, max_logit) = carry - return output.astype(q.dtype), softmax_aux, rng_state + output = output.astype(q.dtype) + # Globalize the rank-local running [H] max across DP/CP. + max_logit = FusedAttnFwdPrimitive._reduce_max_logit_across_mesh( + max_logit, mesh, max_logit_reduce_axes, config + ) + return output, softmax_aux, rng_state, max_logit return mesh, fwd_impl, out_shardings, arg_shardings @@ -3411,6 +3613,7 @@ def fused_attn_fwd( context_parallel_causal_load_balanced: bool = False, context_parallel_axis: str = "", stripe_size: int | None = None, + return_max_logit: bool = False, ) -> jnp.ndarray: """ Perform the forward pass of with cuDNN fused attention implementations. @@ -3450,6 +3653,7 @@ def fused_attn_fwd( Indicates the sequences are ordered for causal mask load balancing when running context parallelism. context_parallel_axis (str): The name of the context parallel axis. stripe_size (int | None): Indicates the striping height to be used for ReorderStrategy.Striped Load Balancing + return_max_logit (bool): Whether to return the per-head maximum attention logit. Returns: (jnp.ndarray): The output tensor from the fused attention. """ @@ -3525,6 +3729,7 @@ def fused_attn_fwd( cp_axis=_maybe_context_parallel_axis(context_parallel_axis), cp_striped_window_size=None, stripe_size=stripe_size, + return_max_logit=return_max_logit, ) primitive = None @@ -3542,7 +3747,7 @@ def fused_attn_fwd( primitive = FusedRingAttnFwdPrimitive.outer_primitive seq_desc_flatten, _ = jax.tree.flatten(sequence_descriptor) - output, softmax_aux, rng_state = primitive.bind( + output, softmax_aux, rng_state, max_logit = primitive.bind( *qkv_for_primitive, bias, softmax_offset, @@ -3551,7 +3756,7 @@ def fused_attn_fwd( config=fused_config, ) rng_state = with_sharding_constraint(rng_state, PartitionSpec(get_all_mesh_axes(), None)) - return (output, softmax_aux, rng_state) + return (output, softmax_aux, rng_state, max_logit) def fused_attn_bwd( diff --git a/transformer_engine/jax/cpp_extensions/base.py b/transformer_engine/jax/cpp_extensions/base.py index 5749a72b0d..aab6173516 100644 --- a/transformer_engine/jax/cpp_extensions/base.py +++ b/transformer_engine/jax/cpp_extensions/base.py @@ -266,10 +266,8 @@ def _gspmd_wrapper(*args, **kwargs): cls.outer_primitive = outer_p -for _name, _value in transformer_engine_jax.registrations().items(): - ffi.register_ffi_target(_name, _value, platform="ROCM" if is_hip_extension() else "CUDA") - # Register EpInstanceState (no-op when TE is built without NCCL EP). +# Custom types must be registered before any FFI handler that references them. if hasattr(transformer_engine_jax, "get_ep_instance_state_type_id"): ffi.register_ffi_type( "EpInstanceState", @@ -281,6 +279,10 @@ def _gspmd_wrapper(*args, **kwargs): ) +for _name, _value in transformer_engine_jax.registrations().items(): + ffi.register_ffi_target(_name, _value, platform="ROCM" if is_hip_extension() else "CUDA") + + def manage_primitives(enable_names=None, disable_names=None, disable_all_first=False): """ Helper function to manage primitive states by name without modifying environment variables. diff --git a/transformer_engine/jax/cpp_extensions/ep.py b/transformer_engine/jax/cpp_extensions/ep.py index 77e60afbcd..ca70ea145c 100644 --- a/transformer_engine/jax/cpp_extensions/ep.py +++ b/transformer_engine/jax/cpp_extensions/ep.py @@ -14,6 +14,7 @@ compound ``(dp, ep)`` axis on the leading dim. """ +import functools from dataclasses import dataclass import jax @@ -24,12 +25,38 @@ import transformer_engine_jax from .base import BasePrimitive, register_primitive from ..sharding import global_mesh_resource, get_mesh_axis_size +from ..version_utils import is_collective_stream_supported + + +def _on_collective_stream(func): + """Pin ``func``'s ops to XLA's collective stream so the scheduler serializes + them with native collectives. No-op on JAX that lacks the annotation.""" + if not is_collective_stream_supported(): + return func + from jax.experimental.compute_on import compute_on + + @functools.wraps(func) + def wrapper(*args, **kwargs): + # compute_on traces its callee and abstract-evals every argument, so it + # cannot take the static EpLayerConfig/PartitionSpec args directly. Wrap + # a nullary thunk that closes over them; the array operands are captured + # as consts and lifted to real operands, outputs stay on device. XLA + # async-wraps the resulting call onto the collective stream. + annotated = compute_on( # pylint: disable=not-callable + compute_type="gpu_stream:collective", + out_memory_spaces=jax.memory.Space.Device, + )(lambda: func(*args, **kwargs)) + return annotated() + + return wrapper + __all__ = [ "EpConfig", "EpLayerConfig", "set_ep_config", "get_ep_config", + "reset_ep_config", "ep_handle_mem_size", "ep_prepare", "ep_dispatch_fwd", @@ -77,6 +104,12 @@ def get_ep_config() -> EpConfig: return _ep_config +def reset_ep_config() -> None: + """Clear the cached EpConfig so a later ep_bootstrap starts fresh (see ep_finalize).""" + global _ep_config + _ep_config = None + + @dataclass(frozen=True) class EpLayerConfig: """Per-layer EP config; mirrors C ``NVTEEpLayerConfig``. @@ -203,8 +236,10 @@ def abstract(topk_idx_aval, *, top_k, dispatch_output_per_expert_alignment, is_o ) leading = _ep_leading_dims(is_outer) token_counts_aval = jax.core.ShapedArray(leading + (num_local_experts,), jnp.int32) + # Per-rank pre-drop recv-slot total (includes tokens dropped on overflow). + total_recv_tokens_aval = jax.core.ShapedArray(leading + (1,), jnp.int32) handle_mem_aval = jax.core.ShapedArray(leading + (handle_mem_size,), jnp.uint8) - return token_counts_aval, handle_mem_aval + return token_counts_aval, total_recv_tokens_aval, handle_mem_aval @staticmethod def outer_abstract(*args, **kwargs): @@ -224,13 +259,13 @@ def lowering(ctx, topk_idx, *, top_k, dispatch_output_per_expert_alignment, is_o @staticmethod def impl(topk_idx, top_k, dispatch_output_per_expert_alignment, is_outer): assert EpPreparePrimitive.inner_primitive is not None - token_counts, handle_mem = EpPreparePrimitive.inner_primitive.bind( + token_counts, total_recv_tokens, handle_mem = EpPreparePrimitive.inner_primitive.bind( topk_idx, top_k=top_k, dispatch_output_per_expert_alignment=dispatch_output_per_expert_alignment, is_outer=is_outer, ) - return token_counts, handle_mem + return token_counts, total_recv_tokens, handle_mem @staticmethod def batcher(batched_args, batch_dims, *, top_k, dispatch_output_per_expert_alignment, is_outer): @@ -250,9 +285,11 @@ def partition( f" with the topk dim replicated; got spec={idx_spec}." ) arg_shardings = tuple(a.sharding for a in arg_infos) - # token_counts / handle_mem inherit the input's leading axis (trailing dims auto-pad to None). + # token_counts / total_recv_tokens / handle_mem inherit the input's leading + # axis (trailing dims auto-pad to None). leading_spec = PartitionSpec(idx_spec[0]) tc_sharding = NamedSharding(mesh, leading_spec) + trt_sharding = NamedSharding(mesh, leading_spec) hm_sharding = NamedSharding(mesh, leading_spec) def sharded_impl(topk_idx): @@ -260,7 +297,7 @@ def sharded_impl(topk_idx): topk_idx, top_k, dispatch_output_per_expert_alignment, False ) - return mesh, sharded_impl, (tc_sharding, hm_sharding), arg_shardings + return mesh, sharded_impl, (tc_sharding, trt_sharding, hm_sharding), arg_shardings @staticmethod def shardy_sharding_rule(*args): @@ -269,7 +306,7 @@ def shardy_sharding_rule(*args): value_types = args[-2] topk_idx_rank = len(value_types[0].shape) in_axes = " ".join(f"L{i}" for i in range(topk_idx_rank - 1)) + " topk" - return f"{in_axes} -> EPL nle, EPL hm" + return f"{in_axes} -> EPL nle, EPL trt, EPL hm" register_primitive(EpPreparePrimitive) @@ -894,8 +931,11 @@ def shardy_sharding_rule(*args): # ── Public-ish helpers (used by jax/ep.py) ────────────────────────────────── +@_on_collective_stream def ep_prepare(cfg: EpLayerConfig, topk_idx): - """Exchange routing metadata for ``cfg``; return ``(token_counts, handle_mem)``.""" + """Exchange routing metadata for ``cfg``; return + ``(token_counts, total_recv_tokens, handle_mem)``. ``total_recv_tokens`` is + the per-rank pre-drop recv-slot total (includes tokens dropped on overflow).""" return EpPreparePrimitive.outer_primitive.bind( topk_idx, top_k=int(cfg.top_k), @@ -904,6 +944,7 @@ def ep_prepare(cfg: EpLayerConfig, topk_idx): ) +@_on_collective_stream def ep_dispatch_fwd( cfg: EpLayerConfig, handle_mem, topk_idx, tokens, topk_weights, recv_capacity_per_rank ): @@ -920,6 +961,7 @@ def ep_dispatch_fwd( ) +@_on_collective_stream def ep_combine_fwd( cfg: EpLayerConfig, handle_mem, expert_out, num_local_tokens, out_partition_spec=None ): @@ -935,6 +977,7 @@ def ep_combine_fwd( ) +@_on_collective_stream def ep_dispatch_bwd( cfg: EpLayerConfig, handle_mem, @@ -956,6 +999,7 @@ def ep_dispatch_bwd( ) +@_on_collective_stream def ep_combine_bwd(cfg: EpLayerConfig, handle_mem, grad, recv_capacity_per_rank): """Backward of combine; returns grad_expert_out [num_procs, recv_capacity_per_rank, H].""" return EpCombineBwdPrimitive.outer_primitive.bind( diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index 64f0289682..3ab2191202 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -991,8 +991,13 @@ def _parse_operand_output_specs( # Non-contracting dims of RHS always needs to be gathered, i.e. for TP + activation_hidden # No batch-dim check needed as `rhs_non_cspecs` never contains batch-dim. # In `rhs_specs`, the batch dim appears only in Wgrad GEMM under `rhs_cspecs`. + # Flatten cspecs since a single element can be a tuple of mesh axes (e.g. ("data", "fsdp")). + flattened_lhs_non_cspecs = [] + for spec in lhs_non_cspecs: + flattened_lhs_non_cspecs.extend(spec if isinstance(spec, tuple) else [spec]) rhs_non_cspecs = tuple( - None if spec in lhs_non_cspecs else spec for spec in rhs_non_cspecs + None if spec in flattened_lhs_non_cspecs or spec in lhs_non_cspecs else spec + for spec in rhs_non_cspecs ) else: @@ -1019,8 +1024,13 @@ def _parse_operand_output_specs( # Non-contracting dims of LHS to be gathered along the SP axis. # Minor note: This causes MaxText TP (= Megatron TP + activation_hidden sharding) gathering x for # dW1 = x^T * dY1 which is unexpected. This is a known issue and no solution has found yet. + # Flatten cspecs since a single element can be a tuple of mesh axes (e.g. ("data", "fsdp")). + flattened_rhs_non_cspecs = [] + for spec in rhs_non_cspecs: + flattened_rhs_non_cspecs.extend(spec if isinstance(spec, tuple) else [spec]) lhs_non_cspecs = tuple( - None if spec in rhs_non_cspecs else spec for spec in lhs_non_cspecs + None if spec in flattened_rhs_non_cspecs or spec in rhs_non_cspecs else spec + for spec in lhs_non_cspecs ) out_specs = lhs_non_cspecs + rhs_non_cspecs diff --git a/transformer_engine/jax/cpp_extensions/quantization.py b/transformer_engine/jax/cpp_extensions/quantization.py index db3950b194..02022fc9c3 100644 --- a/transformer_engine/jax/cpp_extensions/quantization.py +++ b/transformer_engine/jax/cpp_extensions/quantization.py @@ -1299,9 +1299,18 @@ def grouped_quantize( return quantizer.quantize(x, flatten_axis=flatten_axis, group_sizes=group_sizes) n_groups = group_sizes.size original_shape = x.shape - assert n_groups == len( - quantizer.quantizers - ), f"n_groups={n_groups} != n_quantizers = {len(quantizer.quantizers)}" + n_quantizers = len(quantizer.quantizers) + if quantizer.scaling_mode.is_mxfp8_scaling: + # Stateless MXFP8 quantizers may describe a global grouped operation + # while this primitive is traced inside shard_map on only the local + # groups. The recipe is identical for every group, and no per-group + # state is selected here, so the global descriptor only needs to cover + # the local operation. + assert ( + n_groups <= n_quantizers + ), f"local n_groups={n_groups} exceeds global n_quantizers={n_quantizers}" + else: + assert n_groups == n_quantizers, f"n_groups={n_groups} != n_quantizers={n_quantizers}" scale = jnp.ones((n_groups,), jnp.float32) if quantizer.scaling_mode == ScalingMode.DELAYED_TENSOR_SCALING: diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 351647f279..9e4132dad9 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -162,7 +162,7 @@ NVTE_Fused_Attn_Backend GetFusedAttnBackend( NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, float dropout_probability, size_t q_attn_heads, size_t kv_attn_heads, size_t q_max_seqlen, size_t kv_max_seqlen, size_t qk_head_dim, size_t v_head_dim, int64_t window_size_left, - int64_t window_size_right, bool deterministic); + int64_t window_size_right, bool return_max_logit, bool deterministic); pybind11::tuple GetFusedAttnForwardWorkspaceSizes( size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, @@ -170,7 +170,7 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, DType dtype, bool is_training, size_t max_segments_per_seq, int64_t window_size_left, - int64_t window_size_right, bool bottom_right_diagonal); + int64_t window_size_right, bool return_max_logit, bool bottom_right_diagonal); pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, @@ -221,7 +221,8 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedMoEAuxLossBackwardHandler); // the group will dispatch. void SetEpBootstrapParams(pybind11::bytes unique_id_bytes, int ep_size, int rank_within_group, int num_experts, int max_tokens_per_rank, int max_recv_tokens_per_rank, - int hidden_dim, int max_num_sms, int max_token_dtype); + int hidden_dim, int max_num_sms, int max_token_dtype, + bool drop_on_overflow); void ReleaseEpResources(); // Return the handle_mem byte size for a layer config. size_t EpHandleMemSize(int top_k, size_t dispatch_output_per_expert_alignment); diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index ffb15a63c1..5d937dd08c 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -33,12 +33,12 @@ NVTE_Fused_Attn_Backend GetFusedAttnBackend( NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, float dropout_probability, size_t q_attn_heads, size_t kv_attn_heads, size_t q_max_seqlen, size_t kv_max_seqlen, size_t qk_head_dim, size_t v_head_dim, int64_t window_size_left, - int64_t window_size_right, bool deterministic) { + int64_t window_size_right, bool return_max_logit, bool deterministic) { auto backend = nvte_get_fused_attn_backend( is_training, static_cast(q_dtype), static_cast(kv_dtype), qkv_layout, bias_type, mask_type, softmax_type, dropout_probability, q_attn_heads, kv_attn_heads, q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right, - false, false, deterministic); + return_max_logit, false, deterministic); return backend; } @@ -52,8 +52,8 @@ void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t const size_t bias_heads, const size_t q_max_seqlen, const size_t kv_max_seqlen, DType dtype, NVTE_Bias_Type bias_type, NVTE_Fused_Attn_Backend backend, - void *softmax_buf, void *rng_state_buf = nullptr, - void *bias_buf = nullptr, + void *softmax_buf, void *max_logits_buf = nullptr, + void *rng_state_buf = nullptr, void *bias_buf = nullptr, void *softmax_offset_buf = nullptr) { // all backends need softmax but expect different shapes/dtypes tensor_pack->size = 1; @@ -71,13 +71,35 @@ void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t #ifndef USE_ROCM if (backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { #else + // ROCm fused attn has two backends (aotriton and ck); they share the same + // softmax/rng aux tensor shapes and strides, and CK also supports bias, so + // always populate the aux pack for ROCm regardless of backend. { #endif - // ROCm fused attn has two backends: aotriton and ck - // They both have the same shape and stride for softmax and rng aux tensors - // CK now supports bias features - tensor_pack->size = 2; - NVTETensor &rng_state_aux = tensor_pack->tensors[1]; + int size = 1; // Start after softmax. + auto next_aux_tensor = [&]() -> NVTETensor & { + NVTE_CHECK(size < NVTETensorPack::MAX_SIZE, + "Fused attention auxiliary tensor pack capacity exceeded."); + return tensor_pack->tensors[size++]; + }; + +#ifndef USE_ROCM + if (max_logits_buf != nullptr) { + NVTETensor &max_aux = next_aux_tensor(); + NVTEBasicTensor max_aux_data; + max_aux_data.data_ptr = max_logits_buf; + max_aux_data.shape = {}; + max_aux_data.shape.ndim = 4; + max_aux_data.shape.data[0] = input_batch; + max_aux_data.shape.data[1] = attn_heads; + max_aux_data.shape.data[2] = q_max_seqlen; + max_aux_data.shape.data[3] = 1; + max_aux_data.dtype = static_cast(DType::kFloat32); + nvte_set_tensor_param(&max_aux, kNVTERowwiseData, &max_aux_data); + } +#endif + + NVTETensor &rng_state_aux = next_aux_tensor(); NVTEBasicTensor rng_state_aux_data; rng_state_aux_data.data_ptr = rng_state_buf; rng_state_aux_data.shape = {}; @@ -88,12 +110,9 @@ void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t softmax_aux_data.shape.data[3] = 1; // {B,H,Qs,Ks} -> {B,H,Qs,1} softmax_aux_data.dtype = static_cast(DType::kFloat32); - int size = 2; // Start at 2 (we have softmax and rng_state at indices 0, 1) - // include bias if enabled if (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS && bias_type != NVTE_Bias_Type::NVTE_ALIBI) { - NVTETensor &bias_aux = tensor_pack->tensors[size]; - size++; + NVTETensor &bias_aux = next_aux_tensor(); NVTEBasicTensor bias_aux_data; bias_aux_data.data_ptr = bias_buf; bias_aux_data.shape.ndim = 4; @@ -107,8 +126,7 @@ void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t #ifndef USE_ROCM // include softmax_offset if provided if (softmax_offset_buf != nullptr) { - NVTETensor &softmax_offset_aux = tensor_pack->tensors[size]; - size++; + NVTETensor &softmax_offset_aux = next_aux_tensor(); NVTEBasicTensor softmax_offset_aux_data; softmax_offset_aux_data.data_ptr = softmax_offset_buf; softmax_offset_aux_data.shape.ndim = 4; @@ -151,7 +169,7 @@ void PrepareFusedAttnBackwardAuxTensors(NVTETensorPack *tensor_pack, const size_ #endif PrepareFusedAttnForwardAuxTensors(tensor_pack, input_batch, bias_batch, attn_heads, bias_heads, q_max_seqlen, kv_max_seqlen, dtype, dummy_bias_type, - dummy_backend, softmax_buf, rng_state_buf, bias_buf, + dummy_backend, softmax_buf, nullptr, rng_state_buf, bias_buf, softmax_offset_buf); } @@ -161,7 +179,7 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, DType dtype, bool is_training, size_t max_segments_per_seq, int64_t window_size_left, - int64_t window_size_right, bool bottom_right_diagonal) { + int64_t window_size_right, bool return_max_logit, bool bottom_right_diagonal) { auto is_ragged = nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD; auto q_shape = is_ragged ? std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim} @@ -215,8 +233,8 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( dummy_softmax_offset_tensor.data(), s_tensor.data(), o_tensor.data(), &aux_output_tensors, q_cu_seqlens_tensor.data(), kv_cu_seqlens_tensor.data(), ragged_offset_tensor.data(), ragged_offset_tensor.data(), dummy_page_table_tensor.data(), dummy_page_table_tensor.data(), - dummy_rng_state_tensor.data(), q_max_seqlen, kv_max_seqlen, is_training, false, false, - scaling_factor, dropout_probability, qkv_layout, nvte_get_q_format(qkv_layout), + dummy_rng_state_tensor.data(), q_max_seqlen, kv_max_seqlen, is_training, return_max_logit, + false, scaling_factor, dropout_probability, qkv_layout, nvte_get_q_format(qkv_layout), NVTE_QKV_Format_NOT_SET, bias_type, mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, query_workspace_tensor.data(), nullptr); } @@ -260,13 +278,14 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( static void FusedAttnForwardImpl( cudaStream_t stream, void *q, void *k, void *v, void *bias, void *softmax_offset, void *seed, void *q_cu_seqlens, void *kv_cu_seqlens, void *q_seq_offsets, void *k_seq_offsets, void *output, - void *softmax_aux, void *rng_state, void *workspace, size_t input_batch, size_t bias_batch, - size_t q_max_seqlen, size_t kv_max_seqlen, size_t attn_heads, size_t num_gqa_groups, - size_t bias_heads, size_t qk_head_dim, size_t v_head_dim, size_t max_segments_per_seq, - size_t wkspace_size, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, - DType dtype, DType wkspace_dtype, bool is_training, bool deterministic, - int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal) { + void *softmax_aux, void *max_tensor, void *rng_state, void *workspace, size_t input_batch, + size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, size_t attn_heads, + size_t num_gqa_groups, size_t bias_heads, size_t qk_head_dim, size_t v_head_dim, + size_t max_segments_per_seq, size_t wkspace_size, float scaling_factor, + float dropout_probability, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, + NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, DType dtype, DType wkspace_dtype, + bool is_training, bool return_max_logit, bool deterministic, int64_t window_size_left, + int64_t window_size_right, bool bottom_right_diagonal) { FUSED_ATTN_IMPL_COMMON_BLOCK; /* Input tensors */ @@ -281,6 +300,9 @@ static void FusedAttnForwardImpl( // Memset to 0xF0 for filling large negative numbers auto softmax_aux_size = input_batch * q_max_seqlen * attn_heads; (void)cudaMemsetAsync(softmax_aux, 0xF0, softmax_aux_size * sizeof(float), stream); + if (return_max_logit) { + (void)cudaMemsetAsync(max_tensor, 0xF0, softmax_aux_size * sizeof(float), stream); + } } /* Output tensors */ @@ -296,7 +318,7 @@ static void FusedAttnForwardImpl( is_training, static_cast(dtype), static_cast(dtype), qkv_layout, bias_type, mask_type, softmax_type, dropout_probability, attn_heads, num_gqa_groups, q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right, - false, false, deterministic); + return_max_logit, false, deterministic); nvte_populate_rng_state_async(rng_state, seed, q_max_seqlen, kv_max_seqlen, backend, stream); /* Auxiliary tensors (to be propagated to the backward pass later) */ @@ -304,7 +326,8 @@ static void FusedAttnForwardImpl( nvte_tensor_pack_create(&aux_output_tensors); PrepareFusedAttnForwardAuxTensors(&aux_output_tensors, input_batch, bias_batch, attn_heads, bias_heads, q_max_seqlen, kv_max_seqlen, dtype, bias_type, - backend, softmax_aux, softmax_offset); + backend, softmax_aux, return_max_logit ? max_tensor : nullptr, + rng_state, bias, softmax_offset); /* Call the underlying NVTE API */ auto dummy_page_table_tensor = TensorWrapper(nullptr, std::vector{1}, DType::kInt32); @@ -362,7 +385,7 @@ static void FusedAttnForwardImpl( softmax_offset_tensor.data(), s_tensor.data(), o_tensor.data(), &aux_output_tensors, q_cu_seqlens_tensor.data(), kv_cu_seqlens_tensor.data(), q_seq_offsets_tensor.data(), k_seq_offsets_tensor.data(), dummy_page_table_tensor.data(), dummy_page_table_tensor.data(), - rng_state_tensor.data(), q_max_seqlen, kv_max_seqlen, is_training, false, false, + rng_state_tensor.data(), q_max_seqlen, kv_max_seqlen, is_training, return_max_logit, false, scaling_factor, dropout_probability, qkv_layout, nvte_get_q_format(qkv_layout), NVTE_QKV_Format_NOT_SET, bias_type, mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, workspace_tensor.data(), stream); @@ -396,6 +419,7 @@ static void FusedAttnForwardImpl( NVTE_QKV_Layout qkv_layout = \ static_cast(get_attr_value(attrs, "qkv_layout")); \ bool is_training = get_attr_value(attrs, "is_training"); \ + bool return_max_logit = get_attr_value_or_default(attrs, "return_max_logit", false); \ bool deterministic = get_attr_value(attrs, "deterministic"); \ auto is_ragged = nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD; \ size_t wkspace_size = product(workspace_buf->dimensions()); \ @@ -408,8 +432,9 @@ Error_Type FusedAttnForwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Ty Buffer_Type q_cu_seqlens_buf, Buffer_Type kv_cu_seqlens_buf, Buffer_Type q_seq_offsets_buf, Buffer_Type k_seq_offsets_buf, Variadic_Buffer_Type _unused_args, Result_Type output_buf, - Result_Type softmax_aux_buf, Result_Type rng_state_buf, - Result_Type workspace_buf, Dictionary attrs) { + Result_Type softmax_aux_buf, Result_Type max_tensor_buf, + Result_Type rng_state_buf, Result_Type workspace_buf, + Dictionary attrs) { FUSED_ATTN_FFI_GET_ATTRS; FusedAttnForwardImpl( @@ -418,11 +443,12 @@ Error_Type FusedAttnForwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Ty q_cu_seqlens_buf.untyped_data(), kv_cu_seqlens_buf.untyped_data(), is_ragged ? q_seq_offsets_buf.untyped_data() : nullptr, is_ragged ? k_seq_offsets_buf.untyped_data() : nullptr, output_buf->untyped_data(), - softmax_aux_buf->untyped_data(), rng_state_buf->untyped_data(), workspace_buf->untyped_data(), - input_batch, bias_batch, q_max_seqlen, kv_max_seqlen, attn_heads, num_gqa_groups, bias_heads, - qk_head_dim, v_head_dim, max_segments_per_seq, wkspace_size, scaling_factor, - dropout_probability, bias_type, mask_type, softmax_type, qkv_layout, dtype, wkspace_dtype, - is_training, deterministic, window_size_left, window_size_right, bottom_right_diagonal); + softmax_aux_buf->untyped_data(), max_tensor_buf->untyped_data(), + rng_state_buf->untyped_data(), workspace_buf->untyped_data(), input_batch, bias_batch, + q_max_seqlen, kv_max_seqlen, attn_heads, num_gqa_groups, bias_heads, qk_head_dim, v_head_dim, + max_segments_per_seq, wkspace_size, scaling_factor, dropout_probability, bias_type, mask_type, + softmax_type, qkv_layout, dtype, wkspace_dtype, is_training, return_max_logit, deterministic, + window_size_left, window_size_right, bottom_right_diagonal); return ffi_with_cuda_error_check(); } @@ -442,6 +468,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedAttnForwardHandler, FusedAttnForwardFFI, .RemainingArgs() // _cp_aux_args unused .Ret() // output .Ret() // softmax_aux + .Ret() // max_tensor .Ret() // rng_state .Ret() // workspace .Attrs(), diff --git a/transformer_engine/jax/csrc/extensions/ep.cpp b/transformer_engine/jax/csrc/extensions/ep.cpp index cdd7730204..aa5ed27faa 100644 --- a/transformer_engine/jax/csrc/extensions/ep.cpp +++ b/transformer_engine/jax/csrc/extensions/ep.cpp @@ -35,6 +35,7 @@ struct EpBootstrapParams { int hidden_dim = 0; int max_num_sms = 0; NVTEDType max_token_dtype = kNVTEBFloat16; + bool drop_on_overflow = false; }; class EpResources { @@ -53,7 +54,8 @@ class EpResources { .hidden_dim = p.hidden_dim, .num_comm_sms = p.max_num_sms, .max_token_dtype = p.max_token_dtype, - .zero_copy = 0}; + .zero_copy = 0, + .drop_on_overflow = p.drop_on_overflow}; try { nvte_ep_initialize(static_cast(comm_), &cfg); } catch (...) { @@ -125,7 +127,8 @@ struct EpConfig { // synchronize via the UID broadcast). void SetEpBootstrapParams(pybind11::bytes unique_id_bytes_obj, int ep_size, int rank_within_group, int num_experts, int max_tokens_per_rank, int max_recv_tokens_per_rank, - int hidden_dim, int max_num_sms, int max_token_dtype) { + int hidden_dim, int max_num_sms, int max_token_dtype, + bool drop_on_overflow) { std::string uid_str = unique_id_bytes_obj; NVTE_CHECK(static_cast(uid_str.size()) >= 128, "unique_id_bytes must be at least 128 bytes (ncclUniqueId size)."); @@ -143,6 +146,7 @@ void SetEpBootstrapParams(pybind11::bytes unique_id_bytes_obj, int ep_size, int g_ep_params.hidden_dim = hidden_dim; g_ep_params.max_num_sms = max_num_sms; g_ep_params.max_token_dtype = static_cast(max_token_dtype); + g_ep_params.drop_on_overflow = drop_on_overflow; g_ep_params_set = true; } // Acquire outside the lock: EpResources ctor runs ncclCommInitRank which is @@ -196,8 +200,8 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpInstantiateHandler, EpInstantiateImpl, FFI::Bind // ── ep_prepare ──────────────────────────────────────────────────────────────── Error_Type EpPrepareFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_Type topk_idx, - Result_Type recv_tokens_per_expert, Result_Type handle_mem, - EpConfig config) { + Result_Type recv_tokens_per_expert, Result_Type total_recv_tokens, + Result_Type handle_mem, EpConfig config) { (void)ep_state; // lifetime only. auto topk_dims = topk_idx.dimensions(); NVTE_CHECK(topk_dims.size() >= 2, @@ -215,6 +219,10 @@ Error_Type EpPrepareFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_T auto recv_tokens_per_expert_ = TensorWrapper(recv_tokens_per_expert->untyped_data(), tc_shape, DType::kInt32); + std::vector trt_shape = {static_cast(total_recv_tokens->element_count())}; + auto total_recv_tokens_ = + TensorWrapper(total_recv_tokens->untyped_data(), trt_shape, DType::kInt32); + std::vector hm_shape = {static_cast(handle_mem->element_count())}; auto handle_mem_ = TensorWrapper(handle_mem->untyped_data(), hm_shape, DType::kByte); @@ -223,7 +231,7 @@ Error_Type EpPrepareFFI(cudaStream_t stream, EpInstanceState* ep_state, Buffer_T .dispatch_output_per_expert_alignment = static_cast(config.dispatch_output_per_expert_alignment)}; nvte_ep_prepare(handle_mem_.data(), topk_idx_.data(), recv_tokens_per_expert_.data(), - /*total_recv_tokens_per_rank=*/nullptr, &layer_cfg, stream); + total_recv_tokens_.data(), &layer_cfg, stream); return ffi_with_cuda_error_check(); } @@ -233,6 +241,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpPrepareHandler, EpPrepareFFI, .Ctx<::xla::ffi::State>() // EP state .Arg() // topk_idx .Ret() // recv_tokens_per_expert + .Ret() // total_recv_tokens .Ret() // handle_mem .Attrs(), FFI_CudaGraph_Traits); diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index 6cd1c07d23..7672f233aa 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -188,8 +188,8 @@ PYBIND11_MODULE(transformer_engine_jax, m) { m.def("set_ep_bootstrap_params", &SetEpBootstrapParams, pybind11::arg("unique_id_bytes"), pybind11::arg("ep_size"), pybind11::arg("rank_within_group"), pybind11::arg("num_experts"), pybind11::arg("max_tokens_per_rank"), pybind11::arg("max_recv_tokens_per_rank"), - pybind11::arg("hidden_dim"), pybind11::arg("max_num_sms"), - pybind11::arg("max_token_dtype")); + pybind11::arg("hidden_dim"), pybind11::arg("max_num_sms"), pybind11::arg("max_token_dtype"), + pybind11::arg("drop_on_overflow")); m.def("release_ep_resources", &ReleaseEpResources); m.def("ep_handle_mem_size", &EpHandleMemSize, pybind11::arg("top_k"), pybind11::arg("dispatch_output_per_expert_alignment") = 0); diff --git a/transformer_engine/jax/ep.py b/transformer_engine/jax/ep.py index 9704e51546..32074b49e4 100644 --- a/transformer_engine/jax/ep.py +++ b/transformer_engine/jax/ep.py @@ -17,6 +17,8 @@ from transformer_engine.jax.cpp_extensions.ep import _ep_outer_axis from transformer_engine.jax.cpp_extensions.misc import jax_dtype_to_te_dtype from transformer_engine.jax.sharding import ( + _get_mesh, + get_num_devices_in_mesh, global_mesh_resource, get_mesh_axis_size, with_sharding_constraint, @@ -29,6 +31,7 @@ __all__ = [ "EpLayerConfig", "ep_bootstrap", + "ep_finalize", "ep_handle_mem_size", "ep_prepare", "ep_dispatch", @@ -69,6 +72,34 @@ def _allgather_uid(uid_arr, world_size, uid_size): # ── Bootstrap ──────────────────────────────────────────────────────────────── +def _ep_domain_for_rank(mesh, ep_resource, rank, device_to_rank=None): + """Resolve the EP domain (NCCL comm) for ``rank`` from the mesh layout. + + One domain groups ranks sharing all non-ep coordinates, so any orthogonal + axis (tp, pp, cp, ...) yields its own domains. Returns + ``(root_rank, rank_within_group, num_domains)``; ``root_rank`` (ep + coordinate 0) posts the domain's NCCL unique id. + """ + if device_to_rank is None: + + def device_to_rank(d): + return d.process_index + + ep_pos = mesh.axis_names.index(ep_resource) + ep_size = mesh.shape[ep_resource] + ranks = np.vectorize(device_to_rank, otypes=[np.int64])(mesh.devices) + # Move ep last and flatten: each row is one domain (all non-ep coords fixed). + grid = np.moveaxis(ranks, ep_pos, -1).reshape(-1, ep_size) + loc = np.argwhere(grid == rank) + if loc.shape[0] != 1: + raise ValueError( + f"ep_bootstrap: rank {rank} must appear exactly once in the mesh device" + f" grid; found {loc.shape[0]} occurrences." + ) + row, col = int(loc[0, 0]), int(loc[0, 1]) + return int(grid[row, 0]), col, int(grid.shape[0]) + + def ep_bootstrap( world_size, rank, @@ -78,15 +109,17 @@ def ep_bootstrap( hidden_dim, max_token_dtype=jnp.bfloat16, max_num_sms=0, + drop_on_overflow=False, ): """Initialize the EP communicator. Call once per process before any EP op. Must run inside the active JAX Mesh and a global_shard_guard; ep_size and num_ep_groups are read from the mesh axes named by MeshResource.ep_resource - and MeshResource.dp_resource/fsdp_resource. + and MeshResource.dp_resource/fsdp_resource. Axes orthogonal to EP (tp, pp, + cp, ...) are supported and replicated across EP tensors. Args: - world_size: Total number of processes (dp_size * ep_size). + world_size: Total number of processes (product of all mesh axes). rank: Global rank of the calling process. num_experts: Total experts across the EP group. max_tokens_per_rank: Max tokens one rank dispatches per step (sizes send buffers). @@ -95,6 +128,9 @@ def ep_bootstrap( hidden_dim: Feature dimension of token tensors passed to ep_dispatch. max_token_dtype: Widest dtype the group will dispatch (only bfloat16 supported). max_num_sms: SM budget for EP kernels; 0 = auto. + drop_on_overflow: Drop tokens exceeding recv_capacity_per_rank instead of + trapping on overflow. Dropped tokens are still counted in + total_recv_tokens, so callers can detect overflow from it. """ if jnp.dtype(max_token_dtype) != jnp.bfloat16: raise NotImplementedError( @@ -120,30 +156,27 @@ def ep_bootstrap( "ep_bootstrap requires MeshResource.ep_resource to be set; enter a" " global_shard_guard(MeshResource(..., ep_resource=)) before bootstrap." ) - ep_size = get_mesh_axis_size(ep_resource) - outer_axis = _ep_outer_axis() - if outer_axis is None: - if world_size != ep_size: - raise ValueError( - f"ep_bootstrap: world_size ({world_size}) > ep_size ({ep_size}) but neither" - " MeshResource.dp_resource nor fsdp_resource is set; name the outer axis so" - " EP-output tensors can shard across EP groups." - ) - num_ep_groups = 1 - else: - num_ep_groups = get_mesh_axis_size(outer_axis) - if num_ep_groups * ep_size != world_size: + mesh = _get_mesh() + if mesh.empty: raise ValueError( - f"ep_bootstrap: num_ep_groups*ep_size ({num_ep_groups}*{ep_size}=" - f"{num_ep_groups * ep_size}) must equal world_size ({world_size}); check that" - f" the '{outer_axis}' and '{ep_resource}' mesh axes cover all ranks." + "ep_bootstrap must run inside an active jax.sharding.Mesh; enter" + " `with mesh:` (or jax.set_mesh(mesh)) before calling it." ) + if get_num_devices_in_mesh(mesh) != world_size: + raise ValueError( + f"ep_bootstrap: mesh device count ({get_num_devices_in_mesh(mesh)}) must equal" + f" world_size ({world_size})." + ) + ep_size = get_mesh_axis_size(ep_resource) + # num_ep_groups counts only the distinct-token (dp/fsdp) axes; replicated + # axes (tp, pp, ...) do not create distinct EP-output slabs. + outer_axis = _ep_outer_axis() + num_ep_groups = 1 if outer_axis is None else get_mesh_axis_size(outer_axis) if num_experts % ep_size != 0: raise ValueError(f"num_experts ({num_experts}) must be divisible by ep_size ({ep_size}).") UID_SIZE = 128 - dp_color = rank // ep_size - rank_within_group = rank % ep_size + root_rank, rank_within_group, _num_domains = _ep_domain_for_rank(mesh, ep_resource, rank) is_color_root = rank_within_group == 0 if is_color_root: libnccl = ctypes.CDLL("libnccl.so.2", use_errno=True) @@ -156,7 +189,7 @@ def ep_bootstrap( uid_arr = jnp.frombuffer(uid_bytes, dtype=jnp.uint8) all_uids = _allgather_uid(uid_arr, world_size, UID_SIZE) - uid_bytes = bytes(np.asarray(all_uids[dp_color * ep_size]).tolist()) + uid_bytes = bytes(np.asarray(all_uids[root_rank]).tolist()) # Eager NCCL init while ranks are barrier-synced by the UID broadcast above. transformer_engine_jax.set_ep_bootstrap_params( @@ -169,6 +202,7 @@ def ep_bootstrap( hidden_dim, max_num_sms=int(max_num_sms), max_token_dtype=int(jax_dtype_to_te_dtype(max_token_dtype)), + drop_on_overflow=bool(drop_on_overflow), ) # Release the C++ anchor at interpreter shutdown so RAII can tear down NCCL. @@ -192,6 +226,20 @@ def ep_bootstrap( ) +def ep_finalize(): + """Tear down the EP communicator so ``ep_bootstrap`` can run again. + + Only for killing and re-bootstrapping EP mid-program (e.g. tests sweeping + configs); a normal run bootstraps once and lets atexit clean up. Calls the + process-global ``jax.clear_caches()`` so every cached executable releases + the NCCL comm it pins, then frees the EP resources. Call outside any active + EP computation. + """ + jax.clear_caches() + transformer_engine_jax.release_ep_resources() + tex.ep.reset_ep_config() + + def _default_out_partition_spec(): """Leading-axis default: ``(("dp","ep"),)`` if DP/FSDP is set, else ``("ep",)``.""" gsr = global_mesh_resource() @@ -215,8 +263,14 @@ def ep_dispatch(cfg, topk_idx, tokens, topk_weights, recv_capacity_per_rank): ``cfg`` (the pointer-keyed C++ cache keys on handle_mem, not on cfg). Inputs are ``[..., H]`` with only the leading dim sharded as ``ep`` or ``(dp, ep)``. Returns - ``(recv_tokens, recv_topk_weights, handle_mem, token_counts)``; pass - ``handle_mem`` and ``token_counts`` to the matching ``ep_combine``. + ``(recv_tokens, recv_topk_weights, handle_mem, token_counts, total_recv_tokens)``; + pass ``handle_mem`` and ``token_counts`` to the matching ``ep_combine``. + + ``total_recv_tokens`` is the per-rank pre-drop recv-slot count (a + ``[num_procs, 1]`` sharded array); it counts dropped tokens too when + ``drop_on_overflow`` is set. When ``recv_capacity_per_rank`` is not sized for + the worst case, detect overflow by ``process_allgather``-ing it, then + ``max(...) > recv_capacity_per_rank`` flags an overflowing step. """ return _dispatch_fwd(cfg, topk_idx, tokens, topk_weights, recv_capacity_per_rank)[0] @@ -226,12 +280,12 @@ def _dispatch_fwd(cfg, topk_idx, tokens, topk_weights, recv_capacity_per_rank): raise TypeError( f"ep_dispatch: topk_weights must be a floating dtype; got {topk_weights.dtype}." ) - token_counts, handle_mem = tex.ep_prepare(cfg, topk_idx) + token_counts, total_recv_tokens, handle_mem = tex.ep_prepare(cfg, topk_idx) recv_tokens, recv_topk_weights = tex.ep_dispatch_fwd( cfg, handle_mem, topk_idx, tokens, topk_weights, recv_capacity_per_rank ) out_leading = tuple(tokens.shape[:-1]) - primal = (recv_tokens, recv_topk_weights, handle_mem, token_counts) + primal = (recv_tokens, recv_topk_weights, handle_mem, token_counts, total_recv_tokens) return primal, (handle_mem, out_leading) diff --git a/transformer_engine/jax/flax/moe.py b/transformer_engine/jax/flax/moe.py index 3629346e33..cb10c7aa18 100644 --- a/transformer_engine/jax/flax/moe.py +++ b/transformer_engine/jax/flax/moe.py @@ -33,13 +33,11 @@ import jax.numpy as jnp from flax import linen as nn -# Re-exported so downstream users can ``from transformer_engine.jax.flax.moe -# import P`` without a second jax.sharding import. -from jax.sharding import PartitionSpec as P # noqa: F401 # pylint: disable=unused-import - +from transformer_engine.common.recipe import Recipe from ..moe import moe +from ..quantize import QuantizerSet from ..router import ScoreFunction -from ..sharding import get_active_resource_axis +from ..sharding import _get_mesh, get_active_resource_axis from .module import TransformerEngineBase PRNGKey = Any @@ -104,6 +102,9 @@ class _MoEBlock(TransformerEngineBase): If ``True``, multiply expert outputs by their top-k weights *inside* each shard before ``ep_combine`` (saves one global reduction at the cost of an extra broadcast). Default ``False``. + recv_capacity_per_rank : Optional[int] + Exact aligned receive capacity per EP rank. ``None`` reserves the + dropless worst case. The per-expert dispatch-slot alignment is fixed internally at 128 tokens (see ``moe._ALIGN_SIZE``) -- the value required by NCCL EP @@ -117,10 +118,10 @@ class _MoEBlock(TransformerEngineBase): Register per-expert FFN biases (``wi_0_bias``, ``wi_1_bias``, ``wo_bias``). - Quantization is currently configured via the standard TE autocast - context (``fp8_autocast``/``with_quantizer_set``) and threaded - through ``moe()`` internally; this wrapper does not expose a - per-call ``quantizer_sets`` knob yet. + quantization_recipe : Optional[Recipe] + Recipe used to construct the FC1 and FC2 grouped-GEMM quantizer + sets. ``None`` uses the recipe from the active TE autocast context, + or no-op quantizers when autocast is disabled. """ # Architecture @@ -149,6 +150,7 @@ class _MoEBlock(TransformerEngineBase): # MoE knobs forwarded to ``moe()`` apply_topk_weights_early: bool = False + recv_capacity_per_rank: Optional[int] = None # Dtypes / init / misc dtype: DType = jnp.float32 @@ -156,6 +158,7 @@ class _MoEBlock(TransformerEngineBase): bias_init: Initializer = nn.initializers.zeros expert_bias_init: Initializer = nn.initializers.zeros use_ffn_bias: bool = False + quantization_recipe: Optional[Recipe] = None def __post_init__(self): if self.kernel_init is None: @@ -169,7 +172,7 @@ def __post_init__(self): super().__post_init__() @nn.compact - def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array]]: + def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array], Array]: """Run the MoE forward pass. Parameters @@ -184,6 +187,9 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array]]: aux_loss : Optional[jnp.ndarray] Scalar load-balancing loss when ``aux_loss_coeff > 0``, else ``None``. + total_recv_tokens : jnp.ndarray + Non-differentiable per-rank pre-drop recv-slot total; flags + overflow when ``drop_on_overflow`` is set at ep_bootstrap. """ assert ( inputs.ndim == 3 @@ -200,16 +206,14 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array]]: (hidden_size, self.num_experts), self.dtype, ) - wi_0 = self.param( - "wi_0", + # FC1 is stored as one gated-SwiGLU kernel. Keeping its two + # projections contiguous lets the functional MoE path quantize and + # all-gather one FP8 data buffer (and one scale buffer), rather than + # materializing a concatenate inside the custom-VJP. + wi = self.param( + "wi", nn.with_logical_partitioning(self.kernel_init, self.wi_kernel_axes), - (self.num_experts, hidden_size, self.intermediate_size), - self.dtype, - ) - wi_1 = self.param( - "wi_1", - nn.with_logical_partitioning(self.kernel_init, self.wi_kernel_axes), - (self.num_experts, hidden_size, self.intermediate_size), + (self.num_experts, hidden_size, 2 * self.intermediate_size), self.dtype, ) wo = self.param( @@ -250,12 +254,39 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array]]: ) ep_axis = get_active_resource_axis("ep_resource") + mesh = _get_mesh() + data_parallel_size = 1 + for axis in self.data_parallelism_axes: + data_parallel_size *= mesh.shape[axis] + + def make_grouped_quantizer_set(postfix): + # Dispatched token groups span every data-parallel replica, + # whereas expert kernels have one group per global expert. + token_set = self.generate_quantizer_set( + f"{postfix}_token", + fp8_recipe=self.quantization_recipe, + n_groups=data_parallel_size * self.num_experts, + ) + expert_set = self.generate_quantizer_set( + f"{postfix}_expert", + fp8_recipe=self.quantization_recipe, + n_groups=self.num_experts, + ) + return QuantizerSet( + x=token_set.x, + kernel=expert_set.kernel, + dgrad=token_set.dgrad, + ) + + quantizer_sets = ( + make_grouped_quantizer_set("_fc1"), + make_grouped_quantizer_set("_fc2"), + ) return moe( inputs, gate_kernel, - wi_0, - wi_1, + wi, wo, wi_0_bias, wi_1_bias, @@ -271,6 +302,8 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array]]: scaling_factor=self.scaling_factor, aux_loss_coeff=self.aux_loss_coeff, apply_topk_weights_early=self.apply_topk_weights_early, + quantizer_sets=quantizer_sets, + recv_capacity_per_rank=self.recv_capacity_per_rank, ep_axis=ep_axis, data_parallelism_axes=self.data_parallelism_axes, input_axes=self.input_axes, diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 76922d2b55..4b497826cc 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -308,6 +308,7 @@ class _FusedDotProductAttention(nn.Module): # pylint: disable=too-few-public-me score_mod: Optional[Callable] = None score_mod_bprop: Optional[Callable] = None score_mod_requested: bool = False + return_max_logit: bool = False @nn.compact def __call__( @@ -363,6 +364,7 @@ def __call__( "score_mod_bprop": self.score_mod_bprop, "score_mod_tensors": score_mod_tensors, "score_mod_bprop_tensors": score_mod_bprop_tensors, + "return_max_logit": self.return_max_logit, } if self.qkv_layout.is_qkvpacked(): @@ -434,12 +436,17 @@ def __call__( else: raise ValueError(f"Unsupported {self.qkv_layout=}.") + if self.return_max_logit: + x, max_logit = x + if self.transpose_batch_sequence: x = x.transpose([1, 0, 2, 3]) assert ( x.dtype == query.dtype ), f"output dtype {x.dtype} does not match query dtype {query.dtype}" + if self.return_max_logit: + return x, max_logit return x @@ -619,6 +626,9 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods argument to keep tensor operands as normal JAX inputs. score_mod_bprop_tensors: Optional[Mapping[str, Any]], default = None Additional tensors or pass-by-value scalars for ``score_mod_bprop``. + return_max_logit: bool, default = False + If True, return ``(output, max_logit)`` where ``max_logit`` contains the per-head + maximum attention logits with shape ``[h]``. This path requires fused attention. Optimization parameters ----------------------- @@ -647,6 +657,7 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods softmax_type: str = "vanilla" score_mod: Optional[Callable] = None score_mod_bprop: Optional[Callable] = None + return_max_logit: bool = False def __post_init__(self): # TODO(KshitijLakhani): Remove warning in TransformerEngine v2.12 @@ -717,8 +728,8 @@ def __call__( Returns ------- - outputs: jax.numpy.ndarray - Output tensors. + outputs: jax.numpy.ndarray or tuple[jax.numpy.ndarray, jax.numpy.ndarray] + Output tensor, or ``(output, max_logit)`` when ``return_max_logit`` is enabled. """ input_dtype = query.dtype @@ -777,6 +788,8 @@ def __call__( # Use fused attn (if kernel check below passes) by default enable_fused_attn = int(os.getenv("NVTE_FUSED_ATTN", "1")) + if self.return_max_logit and not enable_fused_attn: + raise ValueError("return_max_logit requires fused attention, but NVTE_FUSED_ATTN=0.") sequence_dim = 0 if self.transpose_batch_sequence else 1 seqlen_q = query.shape[sequence_dim] @@ -815,11 +828,17 @@ def __call__( head_dim_qk, head_dim_v, self.window_size, + return_max_logit=self.return_max_logit, ) if score_mod_requested and not has_fused_attn_kernel: raise ValueError( "score_mod requires fused attention, but no fused attention kernel is available." ) + if self.return_max_logit and not has_fused_attn_kernel: + raise ValueError( + "return_max_logit requires fused attention, but no fused attention kernel is " + "available." + ) use_fused_attn = enable_fused_attn and has_fused_attn_kernel @@ -916,6 +935,7 @@ def __call__( score_mod=self.score_mod, score_mod_bprop=self.score_mod_bprop, score_mod_requested=score_mod_requested, + return_max_logit=self.return_max_logit, )( query, key, @@ -927,7 +947,10 @@ def __call__( score_mod_tensors=score_mod_tensors, score_mod_bprop_tensors=score_mod_bprop_tensors, ) - assert x.dtype == input_dtype, f"output_dtype={x.dtype}, input_dtype={input_dtype}" + output = x[0] if self.return_max_logit else x + assert ( + output.dtype == input_dtype + ), f"output_dtype={output.dtype}, input_dtype={input_dtype}" return x diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py index fba3ba4e9b..fa3181dde9 100644 --- a/transformer_engine/jax/moe.py +++ b/transformer_engine/jax/moe.py @@ -22,24 +22,20 @@ ``((*data_parallelism_axes, ep_axis), None, None)``. The public :func:`moe` soft-repins this on entry and warns when a reshard is inserted. -* The EP primitives operate at global view (their custom_partitioning - rules handle per-shard execution). The FFN GEMMs run per-shard inside - a small ``shard_map`` whose ``in_specs`` and ``out_specs`` mirror the - same ``((dp, ep), ...)`` layout. - -Out-of-scope (for now) ----------------------- -FP8 / MXFP8 quantizer sets are not yet wired on this path; turning -them on requires recipe-aware residual specs and ``ScaledTensor`` -leaves across the ``shard_map`` boundary. ``aux_loss_coeff`` and -``expert_bias`` are supported (the former forces a per-step -all-gather over the routing-side logits, which lives off the critical -path and overlaps with the dispatch collective). +* The EP, grouped-quantize, and grouped-GEMM primitives operate at global + view. Their custom partitioning rules handle per-shard execution, + including EP placement and DP/FSDP gathers and reductions. + +FC1 and FC2 use independent quantizer sets. The sets are differentiable +``custom_vjp`` arguments and are returned by the backward rule so +stateful recipes follow the same update semantics as the other TE MLPs. +``aux_loss_coeff`` and ``expert_bias`` are also supported. """ +import math +import warnings from functools import partial from typing import Any, Optional, Tuple, Union -import warnings import flax.struct import jax @@ -48,6 +44,8 @@ from . import cpp_extensions as tex from .quantize import ( + GroupedQuantizer, + QuantizerSet, TensorUsage, noop_quantizer_set, with_sharding_constraint_by_logical_axes, @@ -56,7 +54,7 @@ from .router import ScoreFunction, _validate_score_function from .sharding import _get_mesh -__all__ = ["moe"] +__all__ = ["get_moe_recv_capacity_per_rank", "moe"] # Triton-backed primitives are imported lazily: callers on the PURE_JAX # permutation backend should not need ``triton`` installed. The TRITON @@ -113,6 +111,62 @@ def _require_triton(): _ALIGN_SIZE = 128 +def get_moe_recv_capacity_per_rank( + *, + num_experts: int, + num_experts_per_tok: int, + max_tokens_per_rank: int, + ep_size: int, + recv_capacity_factor: Optional[float] = None, + alignment: int = _ALIGN_SIZE, +) -> int: + """Return the aligned receive capacity for one EP rank. + + ``recv_capacity_factor=None`` reserves the dropless worst case. A finite + factor >= 1 scales the capacity needed by perfectly balanced routing and + is capped at the worst case. The balanced baseline includes the independent + per-local-expert alignment required by NCCL EP. + """ + if num_experts <= 0 or num_experts_per_tok <= 0 or max_tokens_per_rank <= 0: + raise ValueError( + "num_experts, num_experts_per_tok, and max_tokens_per_rank must be positive" + ) + if ep_size <= 0 or num_experts % ep_size != 0: + raise ValueError(f"num_experts={num_experts} must be divisible by ep_size={ep_size}") + if alignment <= 0: + raise ValueError(f"alignment must be positive, got {alignment}") + if recv_capacity_factor is not None: + recv_capacity_factor = float(recv_capacity_factor) + if not math.isfinite(recv_capacity_factor) or recv_capacity_factor < 1.0: + raise ValueError( + "recv_capacity_factor must be finite and >= 1.0, or None for worst-case capacity; " + f"got {recv_capacity_factor}" + ) + + num_local_experts = num_experts // ep_size + tokens_per_ep_group = ep_size * max_tokens_per_rank + max_local_assignments = tokens_per_ep_group * min(num_experts_per_tok, num_local_experts) + max_nonempty_experts = min(num_local_experts, max_local_assignments) + padded_total_bound = max_local_assignments + (alignment - 1) * max_nonempty_experts + aligned_total_bound = ((padded_total_bound + alignment - 1) // alignment) * alignment + per_expert_bound = ( + num_local_experts * ((tokens_per_ep_group + alignment - 1) // alignment) * alignment + ) + worst_case = min(per_expert_bound, aligned_total_bound) + if recv_capacity_factor is None: + return worst_case + + balanced_per_expert = ( + max_tokens_per_rank * num_experts_per_tok + num_local_experts - 1 + ) // num_local_experts + balanced_aligned = ( + num_local_experts * ((balanced_per_expert + alignment - 1) // alignment) * alignment + ) + requested = math.ceil(balanced_aligned * recv_capacity_factor) + requested = ((requested + alignment - 1) // alignment) * alignment + return min(requested, worst_case) + + def _with_sharding_constraint_cast_bwd(x: jnp.ndarray, sharding) -> jnp.ndarray: """Sharding constraint that keeps bwd cotangents in the primal dtype. @@ -158,9 +212,9 @@ def _constraint_bwd(dtype_ref, grad): # cannot run from inside a jit-traced function. The caller must bootstrap # eagerly once per process before any jitted MoE call, then record the # bootstrap signature via ``record_ep_bootstrap_signature_for_moe``. The -# per-call check below verifies the recorded signature is wide enough for -# the current MoE invocation (smaller per-call usage is fine since the C++ -# backend reserves worst-case buffers at bootstrap time). +# per-call check below verifies the recorded signature matches the current +# MoE invocation. NCCL EP permits a smaller token count than the bootstrap +# maximum, but the dispatch receive capacity itself must match exactly. _te_ep_bootstrap_signature: Optional[Tuple[int, int, int, int, int]] = None @@ -193,7 +247,7 @@ def _te_ep_assert_compatible_bootstrap( hidden_dim: int, ep_size: int, ) -> None: - """Verify a prior eager ``ep_bootstrap`` is wide enough for this call.""" + """Verify a prior eager ``ep_bootstrap`` is compatible with this call.""" if _te_ep_bootstrap_signature is None: raise RuntimeError( "TE EP was not bootstrapped. Call" @@ -218,7 +272,7 @@ def _te_ep_assert_compatible_bootstrap( f" (num_experts={num_experts}, max_tokens_per_rank={max_tokens_per_rank}," f" recv_capacity_per_rank={recv_capacity_per_rank}, hidden_dim={hidden_dim}," f" ep_size={ep_size}). Re-bootstrap with wider params (or matching exact" - " sizes) is required." + " sizes) is required. NCCL EP dispatch capacity must exactly match bootstrap." ) @@ -244,7 +298,6 @@ class _Ctx: routing_map: jnp.ndarray cfg: Any = flax.struct.field(pytree_node=False) handle_mem: jnp.ndarray - token_counts: jnp.ndarray recv_topk_weights: jnp.ndarray casted_sorted_x_lhs_trans: Any casted_wi_rhs_trans: Any @@ -254,65 +307,120 @@ class _Ctx: casted_wo_rhs_trans: Any expert_outputs: jnp.ndarray local_group_sizes: jnp.ndarray + quantizer_sets: Any aux_const_buf: Any = None aux_tokens_per_expert: Any = None aux_saved_scores: Any = None # ============================================================================= -# Per-shard FFN body (runs inside shard_map) +# Per-shard FFN body # ============================================================================= +def _validate_moe_quantizer_sets( + quantizer_sets: Tuple[QuantizerSet, QuantizerSet], + *, + num_token_groups: int, + num_expert_groups: int, +) -> None: + """Validate the current global-view MoE quantizer contract. + + Quantizers passed to the public MoE API always describe the global logical + operation. The shard-mapped FFN consumes only its local group count, but it + must not rewrite that public metadata into a shard-local representation. + + Stateful grouped recipes will eventually require sharded leading group + dimensions on their internal state. Until that representation exists, MoE + supports only no-op quantizers and stateless MXFP8 grouped quantizers. + """ + if not isinstance(quantizer_sets, tuple) or len(quantizer_sets) != 2: + raise TypeError("MoE quantizer_sets must be a tuple of FC1 and FC2 QuantizerSet objects.") + + expected_groups = { + "x": num_token_groups, + "kernel": num_expert_groups, + "dgrad": num_token_groups, + } + for set_name, quantizer_set in zip(("FC1", "FC2"), quantizer_sets): + if not isinstance(quantizer_set, QuantizerSet): + raise TypeError(f"MoE {set_name} quantizer must be a QuantizerSet.") + quantizers = { + "x": quantizer_set.x, + "kernel": quantizer_set.kernel, + "dgrad": quantizer_set.dgrad, + } + if all(quantizer is None for quantizer in quantizers.values()): + continue + if any(quantizer is None for quantizer in quantizers.values()): + raise TypeError( + f"MoE {set_name} must use either all no-op quantizers or all grouped MXFP8 " + "quantizers." + ) + + for source, quantizer in quantizers.items(): + if not isinstance(quantizer, GroupedQuantizer): + raise TypeError( + f"MoE {set_name} {source} quantizer must be a GroupedQuantizer; " + f"got {type(quantizer).__name__}." + ) + if not quantizer.scaling_mode.is_mxfp8_scaling: + raise NotImplementedError( + "TE MoE currently supports only BF16/no-op and stateless MXFP8 grouped " + f"quantizers; {set_name} {source} uses {quantizer.scaling_mode}." + ) + if jax.tree_util.tree_leaves(quantizer): + raise NotImplementedError( + "TE MoE does not yet support stateful grouped quantizers. Quantizer state " + "must first be represented with a sharded global group dimension." + ) + expected = expected_groups[source] + if quantizer.n_groups != expected or len(quantizer.quantizers) != expected: + raise ValueError( + f"MoE {set_name} {source} quantizer must describe the global logical " + f"group count {expected}; got n_groups={quantizer.n_groups} and " + f"{len(quantizer.quantizers)} child quantizers." + ) + + def _ffn_fwd_per_shard( recv_tokens_local: jnp.ndarray, recv_topk_weights_local: jnp.ndarray, token_counts_local: jnp.ndarray, - wi_0: jnp.ndarray, - wi_1: jnp.ndarray, + wi: jnp.ndarray, wo: jnp.ndarray, wi_0_bias: Optional[jnp.ndarray], wi_1_bias: Optional[jnp.ndarray], wo_bias: Optional[jnp.ndarray], + quantizer_sets: Tuple[QuantizerSet, QuantizerSet], *, num_local_experts: int, activation_type: str, apply_topk_weights_early: bool, ): - """Per-shard FFN forward. - - Operates on the shard-local ``[1, recv_pr, H]`` slice that - ``tex.ep_dispatch`` produces. Returns the expert outputs (shaped - ``[1, recv_pr, H_out]`` so the surrounding ``shard_map`` reassembles - them as ``[num_procs, recv_pr, H_out]``) plus the residuals consumed - by the bwd. - - ``token_counts_local`` (``[1, num_local_experts]``, from - ``tex.ep_prepare``) is passed to ``grouped_gemm`` as ``group_sizes`` - so cuBLAS skips both 0-token-routed experts and the dispatch - overalloc tail. - """ + """Run the grouped FFN on one shard's EP receive buffer.""" hidden = recv_tokens_local.shape[-1] sorted_x = recv_tokens_local.reshape(-1, hidden) recv_w_flat = recv_topk_weights_local.reshape(-1) - local_group_sizes = token_counts_local.reshape(-1).astype(jnp.int32) + group_sizes = token_counts_local.reshape(-1).astype(jnp.int32) - wi_0 = wi_0.astype(sorted_x.dtype) - wi_1 = wi_1.astype(sorted_x.dtype) + wi = wi.astype(sorted_x.dtype) wo = wo.astype(sorted_x.dtype) - # Concat wi_0/wi_1 along the trailing axis (NOT stack on a new - # axis). grouped_gemm requires the 3D (G, K, N) weight layout with - # contracting_dims=((1,), (1,)); a 4D stack variant walks off the - # end of the RHS and returns NaN. - wi_combined = jnp.concatenate([wi_0, wi_1], axis=-1) + # ``wi`` is stored in its gated-SwiGLU layout [expert, hidden, 2*mlp]. + # Keeping it contiguous lets grouped quantize/GEMM consume it directly. wi_combined_bias = ( jnp.concatenate([wi_0_bias, wi_1_bias], axis=-1) if wi_0_bias is not None else None ) - q_set = noop_quantizer_set - casted_sorted_x = tex.grouped_quantize(sorted_x, q_set.x, local_group_sizes, flatten_axis=-1) - casted_wi = tex.grouped_quantize(wi_combined, q_set.kernel, flatten_axis=-1) + fc1_quantizer_set, fc2_quantizer_set = quantizer_sets + casted_sorted_x = tex.grouped_quantize( + sorted_x, + fc1_quantizer_set.x, + group_sizes, + flatten_axis=-1, + ) + casted_wi = tex.grouped_quantize(wi, fc1_quantizer_set.kernel, flatten_axis=-1) combined_out = tex.grouped_gemm( casted_sorted_x.get_tensor(usage=TensorUsage.LHS), casted_wi.get_tensor(usage=TensorUsage.RHS), @@ -320,8 +428,6 @@ def _ffn_fwd_per_shard( bias=wi_combined_bias, ) gate_proj_out, up_proj_out = jnp.split(combined_out, 2, axis=-1) - casted_sorted_x_lhs_trans = casted_sorted_x.get_tensor(usage=TensorUsage.LHS_TRANS) - casted_wi_rhs_trans = casted_wi.get_tensor(usage=TensorUsage.RHS_TRANS) # Activation inputs (gate_proj_out, up_proj_out) stay in the wi GEMM # output dtype; the activation output (`intermediate`) stays in the @@ -330,45 +436,36 @@ def _ffn_fwd_per_shard( # transitions to the target precision. act_fn = _convert_to_activation_function(activation_type) intermediate = act_fn(gate_proj_out) * up_proj_out - if apply_topk_weights_early: # Fold the per-token combine weights into the FFN intermediate; # the downstream wo GEMM is linear so this is equivalent to the - # late-weighting path. Padded recv slots can contain uninitialized - # data, so overwrite inactive rows with literal zeros instead of - # relying on multiplication by a zero mask (IEEE NaN * 0 = NaN). - # ``w_b`` is cast to ``intermediate.dtype`` so the multiply doesn't - # promote expert_outputs above the EP buffer's element width. - w_b = recv_w_flat[:, None].astype(intermediate.dtype) - active = (recv_w_flat != 0)[:, None] - intermediate = jnp.where(active, intermediate * w_b, jnp.zeros_like(intermediate)) + # late-weighting path. Grouped GEMM skips overallocation tail padding automatically. + # Padding between groups is padded with zeros by NCCL EP. + intermediate = intermediate * recv_w_flat[:, None].astype(intermediate.dtype) casted_intermediate = tex.grouped_quantize( - intermediate, q_set.x, local_group_sizes, flatten_axis=-1 + intermediate, + fc2_quantizer_set.x, + group_sizes, + flatten_axis=-1, ) - casted_wo = tex.grouped_quantize(wo, q_set.kernel, flatten_axis=-1) + casted_wo = tex.grouped_quantize(wo, fc2_quantizer_set.kernel, flatten_axis=-1) expert_outputs = tex.grouped_gemm( casted_intermediate.get_tensor(usage=TensorUsage.LHS), casted_wo.get_tensor(usage=TensorUsage.RHS), contracting_dims=((1,), (1,)), bias=wo_bias, ) - casted_intermediate_lhs_trans = casted_intermediate.get_tensor(usage=TensorUsage.LHS_TRANS) - casted_wo_rhs_trans = casted_wo.get_tensor(usage=TensorUsage.RHS_TRANS) - expert_outputs_3d = expert_outputs.reshape(1, expert_outputs.shape[0], expert_outputs.shape[1]) - # Reshape local_group_sizes to (1, num_local_experts) so the - # surrounding shard_map can stitch per-shard counts back into the - # global (num_procs, num_local_experts) layout matching token_counts. - local_group_sizes_3d = local_group_sizes.reshape(1, num_local_experts) + group_sizes_2d = group_sizes.reshape(1, num_local_experts) residuals = ( - casted_sorted_x_lhs_trans, - casted_wi_rhs_trans, + casted_sorted_x.get_tensor(usage=TensorUsage.LHS_TRANS).checkpoint(fc1_quantizer_set.x), + casted_wi.get_tensor(usage=TensorUsage.RHS_TRANS).checkpoint(fc1_quantizer_set.kernel), gate_proj_out, up_proj_out, - casted_intermediate_lhs_trans, - casted_wo_rhs_trans, - local_group_sizes_3d, + casted_intermediate.get_tensor(usage=TensorUsage.LHS_TRANS).checkpoint(fc2_quantizer_set.x), + casted_wo.get_tensor(usage=TensorUsage.RHS_TRANS).checkpoint(fc2_quantizer_set.kernel), + group_sizes_2d, ) return expert_outputs_3d, residuals @@ -383,28 +480,25 @@ def _ffn_bwd_per_shard( casted_wo_rhs_trans, local_group_sizes: jnp.ndarray, recv_topk_weights_local: jnp.ndarray, + quantizer_sets: Tuple[QuantizerSet, QuantizerSet], *, activation_type: str, apply_topk_weights_early: bool, has_bias: bool, ): - """Per-shard FFN backward. - - Mirrors :func:`_ffn_fwd_per_shard`. Returns - ``(d_sorted_x [1, recv_pr, H], d_recv_w [1, recv_pr], - d_wi_0, d_wi_1, d_wo, d_wi_0_bias, d_wi_1_bias, d_wo_bias)``. - """ - local_group_sizes = local_group_sizes.reshape(-1).astype(jnp.int32) + """Backward mirror of :func:`_ffn_fwd_per_shard`.""" + group_sizes = local_group_sizes.reshape(-1).astype(jnp.int32) d_eo_2d = d_expert_outputs_local.reshape(-1, d_expert_outputs_local.shape[-1]) recv_w_flat = recv_topk_weights_local.reshape(-1) - q_set = noop_quantizer_set - # cuBLAS grouped_gemm skips size_g == 0 groups without zero-filling - # the output slice; mask 0-token-expert wgrads to zero so the - # optimizer never sees uninit memory. - wgrad_group_active = (local_group_sizes > 0)[:, None, None] + fc1_quantizer_set, fc2_quantizer_set = quantizer_sets # wo bwd - casted_d_eo = tex.grouped_quantize(d_eo_2d, q_set.dgrad, local_group_sizes, flatten_axis=-1) + casted_d_eo = tex.grouped_quantize( + d_eo_2d, + fc2_quantizer_set.dgrad, + group_sizes, + flatten_axis=-1, + ) _casted_d_eo_lhs = casted_d_eo.get_tensor(usage=TensorUsage.LHS) _casted_d_eo_rhs = casted_d_eo.get_tensor(usage=TensorUsage.RHS) d_intermediate = tex.grouped_gemm( @@ -417,25 +511,23 @@ def _ffn_bwd_per_shard( _casted_d_eo_rhs, contracting_dims=((0,), (0,)), ) - d_wo = jnp.where(wgrad_group_active, d_wo, jnp.zeros_like(d_wo)) - d_wo_bias = tex.grouped_dbias(d_eo_2d, local_group_sizes) if has_bias else None + d_wo_bias = tex.grouped_dbias(d_eo_2d, group_sizes) if has_bias else None act_fn = _convert_to_activation_function(activation_type) if apply_topk_weights_early: - # intermediate' = intermediate * w * mask. Split the cotangent - # across both factors before the activation bwd consumes it. Padded - # recv slots may still be NaN in the saved activation residuals, so - # use zero-filled residuals on inactive rows before the activation VJP. + # intermediate' = intermediate * w. + # Masking is not required as: + # 1. Padding between groups is zero padded due to NCCL EP. + # 2. Overallocated padding past all groups is uninitialized, but subsequent GEMMs and EP are all group-size aware and will not read past the final group. w_b = recv_w_flat[:, None].astype(d_intermediate.dtype) - active = (recv_w_flat != 0)[:, None] - gate_proj_for_bwd = jnp.where(active, gate_proj_out, jnp.zeros_like(gate_proj_out)) - up_proj_for_bwd = jnp.where(active, up_proj_out, jnp.zeros_like(up_proj_out)) - intermediate_unweighted = act_fn(gate_proj_for_bwd) * up_proj_for_bwd + gate_proj_for_bwd = gate_proj_out + up_proj_for_bwd = up_proj_out + intermediate_unweighted = act_fn(gate_proj_out) * up_proj_out d_recv_w_from_intermediate = jnp.sum( d_intermediate * intermediate_unweighted, axis=-1, ).astype(recv_w_flat.dtype) - d_intermediate = jnp.where(active, d_intermediate * w_b, jnp.zeros_like(d_intermediate)) + d_intermediate = d_intermediate * w_b else: gate_proj_for_bwd = gate_proj_out up_proj_for_bwd = up_proj_out @@ -453,10 +545,13 @@ def _ffn_bwd_per_shard( # gate/up cotangents along the trailing axis, run a single # grouped_quantize + two grouped_gemm pair (one dgrad, one wgrad) # against the fused casted_wi_rhs_trans residual, then split the - # wgrad result back into d_wi_0 / d_wi_1 halves with jnp.split. + # wgrad result remains in the contiguous gated-SwiGLU ``wi`` layout. d_combined = jnp.concatenate([d_gate_proj_out, d_up_proj_out], axis=-1) casted_d_combined = tex.grouped_quantize( - d_combined, q_set.dgrad, local_group_sizes, flatten_axis=-1 + d_combined, + fc1_quantizer_set.dgrad, + group_sizes, + flatten_axis=-1, ) d_sorted_x = tex.grouped_gemm( casted_d_combined.get_tensor(usage=TensorUsage.LHS), @@ -468,10 +563,8 @@ def _ffn_bwd_per_shard( casted_d_combined.get_tensor(usage=TensorUsage.RHS), contracting_dims=((0,), (0,)), ) - d_wi_combined = jnp.where(wgrad_group_active, d_wi_combined, jnp.zeros_like(d_wi_combined)) - d_wi_0, d_wi_1 = jnp.split(d_wi_combined, 2, axis=-1) if has_bias: - d_wi_combined_bias = tex.grouped_dbias(d_combined, local_group_sizes) + d_wi_combined_bias = tex.grouped_dbias(d_combined, group_sizes) d_wi_0_bias, d_wi_1_bias = jnp.split(d_wi_combined_bias, 2, axis=-1) else: d_wi_0_bias = None @@ -482,8 +575,7 @@ def _ffn_bwd_per_shard( return ( d_sorted_x_3d, d_recv_w_3d, - d_wi_0, - d_wi_1, + d_wi_combined, d_wo, d_wi_0_bias, d_wi_1_bias, @@ -499,13 +591,13 @@ def _ffn_bwd_per_shard( def _moe_fwd_rule( x, gate_kernel, - wi_0, - wi_1, + wi, wo, wi_0_bias, wi_1_bias, wo_bias, expert_bias, + quantizer_sets, num_experts, num_experts_per_tok, activation_type, @@ -523,8 +615,9 @@ def _moe_fwd_rule( wo_kernel_axes, dtype, apply_topk_weights_early, + recv_capacity_per_rank, ): - """Forward: gate -> topk -> ep_dispatch -> shard_map(FFN) -> ep_combine. + """Forward: gate -> topk -> ep_dispatch -> FFN -> ep_combine. Returns ``(output, aux_loss)``. ``aux_loss`` is a zero scalar when ``aux_loss_coeff == 0``. @@ -548,6 +641,11 @@ def _moe_fwd_rule( for ax in data_parallelism_axes: dp_size *= mesh.shape[ax] num_procs = num_ep * dp_size + _validate_moe_quantizer_sets( + quantizer_sets, + num_token_groups=dp_size * num_experts, + num_expert_groups=num_experts, + ) B, S, H = x.shape K = num_experts_per_tok @@ -556,18 +654,21 @@ def _moe_fwd_rule( # Per-rank send capacity: B/num_procs rows x S tokens per rank. max_tokens_per_rank = (B // num_procs) * S - # Per-rank receive capacity. NCCL EP HT expert-major lays out variable - # per-expert zones in one flat recv buffer, with each non-empty zone padded - # to ``dispatch_output_per_expert_alignment``. - tokens_per_ep_group = num_ep * max_tokens_per_rank - max_local_assignments = tokens_per_ep_group * min(K, num_local_experts) - max_nonempty_experts = min(num_local_experts, max_local_assignments) - padded_total_bound = max_local_assignments + (_ALIGN_SIZE - 1) * max_nonempty_experts - aligned_total_bound = ((padded_total_bound + _ALIGN_SIZE - 1) // _ALIGN_SIZE) * _ALIGN_SIZE - per_expert_bound = ( - num_local_experts * ((tokens_per_ep_group + _ALIGN_SIZE - 1) // _ALIGN_SIZE) * _ALIGN_SIZE + worst_case_recv_pr = get_moe_recv_capacity_per_rank( + num_experts=num_experts, + num_experts_per_tok=K, + max_tokens_per_rank=max_tokens_per_rank, + ep_size=num_ep, ) - recv_pr = min(per_expert_bound, aligned_total_bound) + if recv_capacity_per_rank is None: + recv_pr = worst_case_recv_pr + else: + recv_pr = int(recv_capacity_per_rank) + if recv_pr <= 0 or recv_pr % _ALIGN_SIZE != 0: + raise ValueError( + f"recv_capacity_per_rank must be a positive multiple of {_ALIGN_SIZE}, got" + f" {recv_pr}" + ) _te_ep_assert_compatible_bootstrap( num_experts=num_experts, @@ -685,7 +786,8 @@ def _moe_fwd_rule( top_k=K, dispatch_output_per_expert_alignment=_ALIGN_SIZE, ) - token_counts, handle_mem = tex.ep_prepare(cfg, topk_idx_3d) + token_counts, total_recv_tokens, handle_mem = tex.ep_prepare(cfg, topk_idx_3d) + token_counts = jax.lax.with_sharding_constraint(token_counts, NamedSharding(mesh, ep2_spec)) recv_tokens, recv_topk_weights = tex.ep_dispatch_fwd( cfg, handle_mem, topk_idx_3d, x, topk_w_3d, recv_pr ) @@ -697,81 +799,57 @@ def _moe_fwd_rule( # ---------------- FFN (per-shard via shard_map) ---------------- has_bias = wi_0_bias is not None kernel_spec = P(ep_axis, None, None) - bias_spec = P(ep_axis, None) if has_bias else None - # token_counts is the per-shard (1, num_local_experts) padded - # per-expert count from ep_prepare; piped into _ffn_fwd_per_shard - # as the grouped_gemm group_sizes so cuBLAS skips both 0-token - # experts and the trailing overalloc tail. - ffn_in_specs = (ep3_spec, ep2_spec, ep2_spec, kernel_spec, kernel_spec, kernel_spec) - ffn_in_args = [recv_tokens, recv_topk_weights, token_counts, wi_0, wi_1, wo] + bias_spec = P(ep_axis, None) + ffn_in_specs = (ep3_spec, ep2_spec, ep2_spec, kernel_spec, kernel_spec) + ffn_in_args = [recv_tokens, recv_topk_weights, token_counts, wi, wo] if has_bias: - ffn_in_specs = ffn_in_specs + (bias_spec, bias_spec, bias_spec) + ffn_in_specs += (bias_spec, bias_spec, bias_spec) ffn_in_args.extend([wi_0_bias, wi_1_bias, wo_bias]) - # FFN residuals live entirely on the local ep rank, so the leading - # "experts" / "rows" dims map to P() (already shard-local). wi is - # fused via jnp.concatenate along the trailing (output) axis - # (see _ffn_fwd_per_shard for rationale), so the residual is a - # single 3D casted_wi_rhs_trans of shape - # (num_local_experts, hidden, 2*H_inter). local_group_sizes is - # now per-shard dynamic (= per-shard token_counts), so its - # residual spec mirrors ep2_spec (one row per ep rank). + # Quantized grouped tensors store their data, scales, and group metadata + # as physical buffers rather than in the source tensor's logical shape. + # A PartitionSpec used as a pytree prefix applies the same ownership to + # every array leaf of the grouped tensor: dispatched-token buffers belong + # to the compound batch shard, while expert-weight buffers belong to EP. + token_buffer_spec = P(batch_pspec_axis) + token_matrix_spec = P(batch_pspec_axis, None) + expert_buffer_spec = P(ep_axis) residuals_spec = ( - P(), # casted_sorted_x_lhs_trans - P(ep_axis, None, None), # casted_wi_rhs_trans - P(), # gate_proj_out - P(), # up_proj_out - P(), # casted_intermediate_lhs_trans - P(ep_axis, None, None), # casted_wo_rhs_trans - ep2_spec, # local_group_sizes (1, num_local_experts) per shard + token_buffer_spec, + expert_buffer_spec, + token_matrix_spec, + token_matrix_spec, + token_buffer_spec, + expert_buffer_spec, + ep2_spec, ) - out_specs = (ep3_spec, residuals_spec) - def _body(*args): + def _ffn_fwd_body(*args): if has_bias: - (r_tok, r_w, tc, w0, w1, w_o, w0b, w1b, wob) = args + r_tok, r_w, tc, local_wi, local_wo, w0b, w1b, wob = args else: - (r_tok, r_w, tc, w0, w1, w_o) = args + r_tok, r_w, tc, local_wi, local_wo = args w0b = w1b = wob = None - # NOTE: tex.ep_dispatch_fwd's NCCL EP HT path leaves the recv - # buffer uninitialised on fully-empty-receiver ranks (and at - # padded slots on partially-loaded ranks). We don't need a - # zero-init guard here anymore because: - # 1. ``tc`` (per-expert padded counts) is plumbed into - # grouped_gemm as group_sizes, so cuBLAS skips both - # 0-token experts and the trailing overalloc tail. - # 2. The per-group wgrad masks in _ffn_bwd_per_shard zero - # ``d_wo`` / ``d_wi_combined`` slices for 0-token-globally - # experts (cuBLAS skips size_g==0 groups without - # zero-filling, which would otherwise leak NaN into the - # user's optimizer). - # 3. All other downstream consumers (ep_combine, - # ep_dispatch_bwd) are handle_mem-aware and read only - # valid positions. - # If a future caller adds a non-group-aware reader of r_tok - # (e.g. an inspect probe over the full recv tile), re-add the - # ``jax.lax.cond(jnp.any(r_w != 0), identity, zeros_like)`` - # guard here. return _ffn_fwd_per_shard( r_tok, r_w, tc, - w0, - w1, - w_o, + local_wi, + local_wo, w0b, w1b, wob, + quantizer_sets, num_local_experts=num_local_experts, activation_type=activation_type, apply_topk_weights_early=apply_topk_weights_early, ) expert_outputs, ffn_residuals = shard_map( - _body, + _ffn_fwd_body, mesh=mesh, in_specs=ffn_in_specs, - out_specs=out_specs, + out_specs=(ep3_spec, residuals_spec), check_rep=False, )(*ffn_in_args) expert_outputs = jax.lax.with_sharding_constraint(expert_outputs, NamedSharding(mesh, ep3_spec)) @@ -799,6 +877,8 @@ def _body(*args): num_local_tokens=(B, S), out_partition_spec=out_partition_spec, ) + # output of MLP should be sharded the same way as the activation input + output = with_sharding_constraint_by_logical_axes(output, input_axes) ( casted_sorted_x_lhs_trans, @@ -819,7 +899,6 @@ def _body(*args): routing_map=routing_map, cfg=cfg, handle_mem=handle_mem, - token_counts=token_counts, recv_topk_weights=recv_topk_weights, casted_sorted_x_lhs_trans=casted_sorted_x_lhs_trans, casted_wi_rhs_trans=casted_wi_rhs_trans, @@ -829,6 +908,7 @@ def _body(*args): casted_wo_rhs_trans=casted_wo_rhs_trans, expert_outputs=expert_outputs, local_group_sizes=local_group_sizes, + quantizer_sets=quantizer_sets, aux_const_buf=aux_const_buf, aux_tokens_per_expert=aux_tokens_per_expert, aux_saved_scores=aux_saved_scores, @@ -838,7 +918,8 @@ def _body(*args): "x_shape": x.shape, "recv_pr": recv_pr, } - return (output, aux_loss), (ctx, static) + # total_recv_tokens is a non-differentiable overflow signal (see moe()). + return (output, aux_loss, total_recv_tokens), (ctx, static) def _moe_bwd_rule( @@ -859,14 +940,16 @@ def _moe_bwd_rule( wo_kernel_axes, dtype, apply_topk_weights_early, + recv_capacity_per_rank, residuals, cotangents, ): """Backward mirror of :func:`_moe_fwd_rule`.""" - del num_groups, group_topk, dtype # captured in residuals / unused in bwd + del num_groups, group_topk, dtype, recv_capacity_per_rank # captured / unused in bwd from jax.experimental.shard_map import shard_map - d_output, d_aux_loss = cotangents + # total_recv_tokens is a non-differentiable output; its cotangent is unused. + d_output, d_aux_loss, _d_total_recv_tokens = cotangents ctx, static = residuals has_bias = static["has_bias"] @@ -876,10 +959,6 @@ def _moe_bwd_rule( mesh = _get_mesh() if mesh is None or mesh.empty: raise ValueError("moe(...) requires an active jax.sharding.Mesh.") - dp_size = 1 - for ax in data_parallelism_axes: - dp_size *= mesh.shape[ax] - B, S, _ = x_shape K = num_experts_per_tok if not data_parallelism_axes: @@ -896,39 +975,33 @@ def _moe_bwd_rule( grad_pre_combine = jax.lax.with_sharding_constraint( grad_pre_combine, NamedSharding(mesh, ep3_spec) ) - if apply_topk_weights_early: # combine_fwd consumed already-weighted expert_outputs; the recv_w # cotangent flows through the early-weighting step inside the FFN bwd. d_expert_outputs = grad_pre_combine d_recv_w_from_combine = jnp.zeros_like(ctx.recv_topk_weights) else: - # Reverse the late-weighting multiply. Padded expert-major rows are - # part of the physical grouped-GEMM ranges, so write literal zero - # cotangents for inactive rows instead of relying on NaN * 0. w = ctx.recv_topk_weights[..., None].astype(grad_pre_combine.dtype) - mask_bool = (ctx.recv_topk_weights != 0)[..., None] - d_expert_outputs = jnp.where( - mask_bool, grad_pre_combine * w, jnp.zeros_like(grad_pre_combine) - ) + d_expert_outputs = grad_pre_combine * w d_recv_w_from_combine = (grad_pre_combine * ctx.expert_outputs).sum(axis=-1) d_recv_w_from_combine = d_recv_w_from_combine.astype(ctx.recv_topk_weights.dtype) # ---------------- FFN bwd (per-shard via shard_map) ---------------- kernel_spec = P(ep_axis, None, None) - bias_spec = P(ep_axis, None) if has_bias else None - - bwd_in_specs = ( - ep3_spec, # d_expert_outputs - P(), # casted_sorted_x_lhs_trans - P(ep_axis, None, None), # casted_wi_rhs_trans - P(), # gate_proj_out - P(), # up_proj_out - P(), # casted_intermediate_lhs_trans - P(ep_axis, None, None), # casted_wo_rhs_trans - ep2_spec, # local_group_sizes (1, num_local_experts) per shard - ep2_spec, # recv_topk_weights + bias_spec = P(ep_axis, None) + token_buffer_spec = P(batch_pspec_axis) + token_matrix_spec = P(batch_pspec_axis, None) + expert_buffer_spec = P(ep_axis) + residuals_specs = ( + token_buffer_spec, + expert_buffer_spec, + token_matrix_spec, + token_matrix_spec, + token_buffer_spec, + expert_buffer_spec, + ep2_spec, ) + bwd_in_specs = (ep3_spec, *residuals_specs, ep2_spec) bwd_in_args = [ d_expert_outputs, ctx.casted_sorted_x_lhs_trans, @@ -940,67 +1013,65 @@ def _moe_bwd_rule( ctx.local_group_sizes, ctx.recv_topk_weights, ] - bwd_out_specs = ( - ep3_spec, # d_sorted_x - ep2_spec, # d_recv_w_from_intermediate - kernel_spec, # d_wi_0 - kernel_spec, # d_wi_1 - kernel_spec, # d_wo - bias_spec if has_bias else None, # d_wi_0_bias - bias_spec if has_bias else None, # d_wi_1_bias - bias_spec if has_bias else None, # d_wo_bias - ) - def _bwd_body(*args): - ( - d_sorted_x_3d, - d_recv_w_3d, - d_wi_0, - d_wi_1, - d_wo, - d_wi_0_bias, - d_wi_1_bias, - d_wo_bias, - ) = _ffn_bwd_per_shard( + def _ffn_bwd_body(*args): + grads = _ffn_bwd_per_shard( *args, + ctx.quantizer_sets, activation_type=activation_type, apply_topk_weights_early=apply_topk_weights_early, has_bias=has_bias, ) - # Weight grads accumulate per-DP-shard inside the body; psum across - # DP axes so each replica sees the full sum (matches out_specs - # P(ep_axis, ...) which is DP-replicated). + ( + d_sorted_x_local, + d_recv_w_local, + d_wi_local, + d_wo_local, + d_wi_0_bias_local, + d_wi_1_bias_local, + d_wo_bias_local, + ) = grads if data_parallelism_axes: - dp = tuple(data_parallelism_axes) - d_wi_0 = jax.lax.psum(d_wi_0, axis_name=dp) - d_wi_1 = jax.lax.psum(d_wi_1, axis_name=dp) - d_wo = jax.lax.psum(d_wo, axis_name=dp) + dp_axes = tuple(data_parallelism_axes) + d_wi_local = jax.lax.psum(d_wi_local, axis_name=dp_axes) + d_wo_local = jax.lax.psum(d_wo_local, axis_name=dp_axes) if has_bias: - d_wi_0_bias = jax.lax.psum(d_wi_0_bias, axis_name=dp) - d_wi_1_bias = jax.lax.psum(d_wi_1_bias, axis_name=dp) - d_wo_bias = jax.lax.psum(d_wo_bias, axis_name=dp) + d_wi_0_bias_local = jax.lax.psum(d_wi_0_bias_local, axis_name=dp_axes) + d_wi_1_bias_local = jax.lax.psum(d_wi_1_bias_local, axis_name=dp_axes) + d_wo_bias_local = jax.lax.psum(d_wo_bias_local, axis_name=dp_axes) return ( - d_sorted_x_3d, - d_recv_w_3d, - d_wi_0, - d_wi_1, - d_wo, - d_wi_0_bias, - d_wi_1_bias, - d_wo_bias, + d_sorted_x_local, + d_recv_w_local, + d_wi_local, + d_wo_local, + d_wi_0_bias_local, + d_wi_1_bias_local, + d_wo_bias_local, ) + if has_bias: + bwd_out_specs = ( + ep3_spec, + ep2_spec, + kernel_spec, + kernel_spec, + bias_spec, + bias_spec, + bias_spec, + ) + else: + bwd_out_specs = (ep3_spec, ep2_spec, kernel_spec, kernel_spec, None, None, None) + ( d_sorted_x, d_recv_w_from_intermediate, - d_wi_0, - d_wi_1, + d_wi, d_wo, d_wi_0_bias, d_wi_1_bias, d_wo_bias, ) = shard_map( - _bwd_body, + _ffn_bwd_body, mesh=mesh, in_specs=bwd_in_specs, out_specs=bwd_out_specs, @@ -1084,9 +1155,14 @@ def _bwd_body(*args): # optimizers see consistent shardings. d_x = with_sharding_constraint_by_logical_axes(d_x, input_axes) d_gate_kernel = with_sharding_constraint_by_logical_axes(d_gate_kernel, gate_kernel_axes) - d_wi_0 = with_sharding_constraint_by_logical_axes(d_wi_0, wi_kernel_axes) - d_wi_1 = with_sharding_constraint_by_logical_axes(d_wi_1, wi_kernel_axes) + d_wi = with_sharding_constraint_by_logical_axes(d_wi, wi_kernel_axes) d_wo = with_sharding_constraint_by_logical_axes(d_wo, wo_kernel_axes) + if has_bias: + wi_bias_axes = (wi_kernel_axes[0], *wi_kernel_axes[2:]) + wo_bias_axes = (wo_kernel_axes[0], *wo_kernel_axes[2:]) + d_wi_0_bias = with_sharding_constraint_by_logical_axes(d_wi_0_bias, wi_bias_axes) + d_wi_1_bias = with_sharding_constraint_by_logical_axes(d_wi_1_bias, wi_bias_axes) + d_wo_bias = with_sharding_constraint_by_logical_axes(d_wo_bias, wo_bias_axes) # expert_bias has no learnable bwd path through fused_topk: the # primitive's bwd returns None for the bias slot. Match that with a @@ -1097,13 +1173,13 @@ def _bwd_body(*args): return ( d_x, d_gate_kernel, - d_wi_0, - d_wi_1, + d_wi, d_wo, d_wi_0_bias if has_bias else None, d_wi_1_bias if has_bias else None, d_wo_bias if has_bias else None, d_expert_bias, + ctx.quantizer_sets, ) @@ -1112,17 +1188,17 @@ def _bwd_body(*args): # ============================================================================= -@partial(jax.custom_vjp, nondiff_argnums=tuple(range(9, 26))) +@partial(jax.custom_vjp, nondiff_argnums=tuple(range(9, 27))) def _moe( x, gate_kernel, - wi_0, - wi_1, + wi, wo, wi_0_bias, wi_1_bias, wo_bias, expert_bias, + quantizer_sets, num_experts, num_experts_per_tok, activation_type, @@ -1140,17 +1216,18 @@ def _moe( wo_kernel_axes, dtype, apply_topk_weights_early, + recv_capacity_per_rank, ): primal, _ = _moe_fwd_rule( x, gate_kernel, - wi_0, - wi_1, + wi, wo, wi_0_bias, wi_1_bias, wo_bias, expert_bias, + quantizer_sets, num_experts, num_experts_per_tok, activation_type, @@ -1168,6 +1245,7 @@ def _moe( wo_kernel_axes, dtype, apply_topk_weights_early, + recv_capacity_per_rank, ) return primal @@ -1178,8 +1256,7 @@ def _moe( def moe( x: jnp.ndarray, gate_kernel: jnp.ndarray, - wi_0: jnp.ndarray, - wi_1: jnp.ndarray, + wi: jnp.ndarray, wo: jnp.ndarray, wi_0_bias: Optional[jnp.ndarray] = None, wi_1_bias: Optional[jnp.ndarray] = None, @@ -1196,6 +1273,10 @@ def moe( scaling_factor: float = 1.0, aux_loss_coeff: float = 0.0, apply_topk_weights_early: bool = False, + quantizer_sets: Tuple[QuantizerSet, QuantizerSet] = ( + noop_quantizer_set, + noop_quantizer_set, + ), ep_axis: str, data_parallelism_axes: Tuple[str, ...] = (), input_axes: Tuple[Optional[str], ...] = (), @@ -1203,11 +1284,14 @@ def moe( wi_kernel_axes: Tuple[Optional[str], ...] = ("exp", "embed", "mlp"), wo_kernel_axes: Tuple[Optional[str], ...] = ("exp", "mlp", "embed"), dtype: jnp.dtype = jnp.float32, -) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: + recv_capacity_per_rank: Optional[int] = None, +) -> Tuple[jnp.ndarray, Optional[jnp.ndarray], jnp.ndarray]: """Run a full MoE block under a single fused custom_vjp on the TE EP path. - Returns ``(output, aux_loss)``. ``aux_loss`` is ``None`` when - ``aux_loss_coeff == 0`` and a 0-d scalar otherwise. + Returns ``(output, aux_loss, total_recv_tokens)``. ``aux_loss`` is ``None`` + when ``aux_loss_coeff == 0``, else a 0-d scalar. ``total_recv_tokens`` is a + non-differentiable pre-drop recv-slot total (grad ``None``); see + ``ep_dispatch`` for using it to detect overflow. Parameters ---------- @@ -1223,6 +1307,18 @@ def moe( all-gather over the routing-side logits is inserted so the ``fused_moe_aux_loss`` kernel sees a global ``[T_global, E]`` view; this lives off the dispatch critical path. + quantizer_sets : Tuple[QuantizerSet, QuantizerSet] + Independent FC1 and FC2 quantizer sets describing the global logical + operation. Token quantizers have ``dp_size * num_experts`` groups and + kernel quantizers have ``num_experts`` groups; shard-local FFN calls use + this global descriptor unchanged. Currently only no-op (BF16) and + stateless grouped MXFP8 quantizers are supported. They are differentiable + custom-VJP arguments so recipe state is threaded through backward. + recv_capacity_per_rank : Optional[int] + Exact aligned receive-buffer capacity for each EP rank. ``None`` + (default) reserves the dropless aligned worst case. The value must match + the capacity used by ``ep_bootstrap``. Overflow is reported through + ``total_recv_tokens`` when bootstrap used ``drop_on_overflow=True``. Note that the per-expert dispatch-slot alignment is fixed internally at 128 tokens (``_ALIGN_SIZE``); see that constant's docstring for @@ -1233,7 +1329,7 @@ def moe( * ``ep_axis`` and ``data_parallelism_axes`` are *physical mesh axis names* -- they index ``jax.sharding.Mesh.shape`` directly (to compute ``num_ep`` / ``dp_size`` and to construct - ``P((dp..., ep), None, None)`` for the per-shard + ``P((dp..., ep), None, None)`` for the physical ``jax.lax.with_sharding_constraint`` calls that JAX requires to refer to real mesh axes). * ``input_axes``, ``gate_kernel_axes``, ``wi_kernel_axes``, @@ -1286,16 +1382,16 @@ def moe( else: expert_bias_arg = expert_bias.astype(jnp.float32) - output, aux_loss = _moe( + output, aux_loss, total_recv_tokens = _moe( x, gate_kernel, - wi_0, - wi_1, + wi, wo, wi_0_bias, wi_1_bias, wo_bias, expert_bias_arg, + quantizer_sets, num_experts, num_experts_per_tok, activation_type, @@ -1313,8 +1409,9 @@ def moe( wo_kernel_axes, dtype, apply_topk_weights_early, + recv_capacity_per_rank, ) if aux_loss_coeff <= 0.0: aux_loss = None assert output.dtype == x.dtype, f"moe() output dtype {output.dtype} != input dtype {x.dtype}" - return output, aux_loss + return output, aux_loss, total_recv_tokens diff --git a/transformer_engine/jax/quantize/tensor.py b/transformer_engine/jax/quantize/tensor.py index c5ad0451fd..edcec01924 100644 --- a/transformer_engine/jax/quantize/tensor.py +++ b/transformer_engine/jax/quantize/tensor.py @@ -8,9 +8,10 @@ both single-scale (1x) and double-scale (2x) quantization schemes. It supports rowwise and colwise quantization modes with proper scaling and dequantization. """ +import math +from abc import ABC, abstractmethod from dataclasses import dataclass from typing import Callable, Optional, Tuple -from abc import ABC, abstractmethod import jax.numpy as jnp from jax.tree_util import register_pytree_node_class @@ -436,6 +437,15 @@ def __post_init__(self): 0 < self.flatten_axis < data_ndim ), f"flatten_axis {self.flatten_axis} is out of bounds for data.ndim = {data_ndim}" + # A grouped tensor used as a shard_map residual temporarily has global + # physical data/scale buffers while ``original_shape`` intentionally + # continues to describe the shard-local logical tensor. The matching + # shard_map input restores local leaves before grouped GEMM consumes + # the wrapper. Validate scale layout only when this is a genuine local + # tensor view; the global transport view is not itself a GEMM operand. + if self.data.size != math.prod(self.original_shape): + return + active_dims = ( self.first_dims if self.first_dims is not None and self.first_dims.size > 0 diff --git a/transformer_engine/jax/router.py b/transformer_engine/jax/router.py index 80ec42a95f..170c74fe5f 100644 --- a/transformer_engine/jax/router.py +++ b/transformer_engine/jax/router.py @@ -342,12 +342,11 @@ def _fused_moe_aux_loss_fwd(probs, tokens_per_expert, topk, coeff): def _fused_moe_aux_loss_bwd(topk, coeff, residuals, g): del topk, coeff const_buf, tokens_per_expert, num_tokens = residuals - grad_aux_loss = g.reshape(1) grad_probs = fused_moe_aux_loss_bwd( const_buf, tokens_per_expert, - grad_aux_loss, + g, num_tokens, ) return grad_probs, None diff --git a/transformer_engine/jax/version_utils.py b/transformer_engine/jax/version_utils.py index e4619d8670..500c859b4c 100644 --- a/transformer_engine/jax/version_utils.py +++ b/transformer_engine/jax/version_utils.py @@ -64,6 +64,26 @@ def is_triton_autotuned_alias_safe() -> bool: return v >= PkgVersion(_TRITON_AUTOTUNED_ALIAS_STABLE_FLOOR) +# XLA gained the ``gpu_stream:collective`` stream annotation in openxla/xla#39604, +# first shipping in JAX 0.10.1. Older XLA fatally fails on it. +# However, JAX 0.10.1 renamed this API to compute_on2 and slightly changed the +# signature, then 0.11.1 renamed it back to compute_on. For simplicity, +# we will support 0.11.1+ +_COLLECTIVE_STREAM_MIN_JAX_VERSION = "0.11.1" + + +@lru_cache(maxsize=None) +def is_collective_stream_supported() -> bool: + """Return True if the installed JAX supports the gpu_stream:collective annotation.""" + if not jax_version_meet_requirement(_COLLECTIVE_STREAM_MIN_JAX_VERSION): + return False + try: + from jax.experimental.compute_on import compute_on # pylint: disable=unused-import + except ImportError: + return False + return True + + def is_triton_extension_supported() -> bool: """Return True if the current JAX version supports Triton kernel dispatch. @@ -77,6 +97,7 @@ def is_triton_extension_supported() -> bool: __all__ = [ "jax_version_meet_requirement", "is_triton_autotuned_alias_safe", + "is_collective_stream_supported", "is_triton_extension_supported", "TRITON_EXTENSION_MIN_JAX_VERSION", "TRITON_EXTENSION_CUDA_GRAPH_MIN_JAX_VERSION", diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 0d578b0300..b4698fab16 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -31,6 +31,7 @@ from transformer_engine.pytorch.module import destroy_ub from transformer_engine.pytorch.module import UserBufferQuantizationMode from transformer_engine.pytorch.attention import DotProductAttention +from transformer_engine.pytorch.attention import FusedMLAQUpProjRopeQuant from transformer_engine.pytorch.attention import MultiheadAttention from transformer_engine.pytorch.attention import InferenceParams from transformer_engine.pytorch.attention import RotaryPositionEmbedding @@ -75,9 +76,10 @@ from transformer_engine.pytorch.cross_entropy import parallel_cross_entropy from torch.utils.cpp_extension import IS_HIP_EXTENSION as _IS_HIP_EXTENSION if not _IS_HIP_EXTENSION: - from transformer_engine.pytorch.newton_schulz import ( + from transformer_engine.pytorch.optimizers.newton_schulz import ( CusolverMpCtx, newton_schulz, + newton_schulz_tp, ) from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage from transformer_engine.pytorch.quantized_tensor import QuantizedTensor @@ -98,6 +100,12 @@ from transformer_engine.pytorch.tensor import MXFP8Tensor from transformer_engine.pytorch.tensor import Float8BlockwiseQTensor from transformer_engine.pytorch.tensor import NVFP4Tensor +from transformer_engine.pytorch.tensor import HybridQuantizer +from transformer_engine.pytorch.tensor import HybridQuantizedTensorStorage +from transformer_engine.pytorch.tensor import IdentityQuantizer +from transformer_engine.pytorch.tensor import IdentityTensorStorage +from transformer_engine.pytorch.tensor import HybridQuantizedTensor +from transformer_engine.pytorch.tensor import IdentityTensor from transformer_engine.pytorch.tensor.float8_tensor import ( _make_float8_tensor_in_reduce_ex, ) @@ -110,6 +118,12 @@ from transformer_engine.pytorch.tensor.float8_blockwise_tensor import ( _make_float8_blockwise_tensor_in_reduce_ex, ) +from transformer_engine.pytorch.tensor.hybrid_tensor import ( + _make_hybrid_quantized_tensor_in_reduce_ex, +) +from transformer_engine.pytorch.tensor.identity_tensor import ( + _make_identity_tensor_in_reduce_ex, +) try: torch._dynamo.config.error_on_nested_jit_trace = False @@ -134,6 +148,8 @@ MXFP8TensorStorage, NVFP4TensorStorage, Float8BlockwiseQTensorStorage, + HybridQuantizedTensorStorage, + IdentityTensorStorage, # Quantizer types embedded in metadata Quantizer, Float8Quantizer, @@ -141,6 +157,8 @@ MXFP8Quantizer, NVFP4Quantizer, Float8BlockQuantizer, + HybridQuantizer, + IdentityQuantizer, # Python IntEnum used as Quantizer.dtype. DType, # pybind11 enum used as Quantizer.dtype. @@ -151,6 +169,8 @@ _make_mxfp8_tensor_in_reduce_ex, _make_nvfp4_tensor_in_reduce_ex, _make_float8_blockwise_tensor_in_reduce_ex, + _make_hybrid_quantized_tensor_in_reduce_ex, + _make_identity_tensor_in_reduce_ex, ] ) except (ImportError, AttributeError): diff --git a/transformer_engine/pytorch/attention/__init__.py b/transformer_engine/pytorch/attention/__init__.py index c4c2aa3e72..f6e4f0b37f 100644 --- a/transformer_engine/pytorch/attention/__init__.py +++ b/transformer_engine/pytorch/attention/__init__.py @@ -5,12 +5,14 @@ """Python interface for attention""" from .dot_product_attention import DotProductAttention +from .fused_mla_q_uproj import FusedMLAQUpProjRopeQuant from .multi_head_attention import MultiheadAttention from .inference import InferenceParams from .rope import RotaryPositionEmbedding __all__ = [ "DotProductAttention", + "FusedMLAQUpProjRopeQuant", "MultiheadAttention", "InferenceParams", "RotaryPositionEmbedding", diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 83ab12c364..4fe6f28773 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -5,6 +5,7 @@ # See LICENSE for license information. """Attention Backends.""" + from contextlib import nullcontext from importlib.metadata import version as get_pkg_version from importlib.metadata import PackageNotFoundError @@ -129,10 +130,10 @@ except PackageNotFoundError: pass # only print warning if use_flash_attention_2 = True in get_attention_backend else: - if torch.cuda.is_available() and ( (not IS_HIP_EXTENSION) and get_device_compute_capability() >= (10, 0)): - if fa_utils.version_required_blackwell <= fa_utils.version <= fa_utils.max_version: + if torch.cuda.is_available() and ((not IS_HIP_EXTENSION) and get_device_compute_capability() >= (10, 0)): + if fa_utils.is_version_supported(fa_utils.version, fa_utils.version_required_blackwell): fa_utils.is_installed = True - elif fa_utils.version_required <= fa_utils.version <= fa_utils.max_version: + elif fa_utils.is_version_supported(fa_utils.version, fa_utils.version_required): fa_utils.is_installed = True if fa_utils.is_installed: @@ -200,14 +201,34 @@ flash_attn_func_v4 = None flash_attn_varlen_func_v4 = None else: - from flash_attn.cute.interface import ( # pylint: disable=ungrouped-imports,no-name-in-module - flash_attn_func as flash_attn_func_v4, - flash_attn_varlen_func as flash_attn_varlen_func_v4, - _validate_head_dims as _fa4_validate_head_dims, - ) + try: + cutlass_dsl_version = PkgVersion(get_pkg_version("nvidia-cutlass-dsl")) + + # FA4 4.0.0b24 requires CUTLASS DSL 4.6.2 or newer. + if fa_utils.fa4_version == PkgVersion("4.0.0b24") and cutlass_dsl_version < PkgVersion( + "4.6.2" + ): + raise ImportError( + "flash-attn-4 4.0.0b24 requires nvidia-cutlass-dsl>=4.6.2; " + f"found {cutlass_dsl_version}" + ) - fa_utils.v4_validate_head_dims = _fa4_validate_head_dims - fa_utils.set_flash_attention_4_params() + from flash_attn.cute.interface import ( # pylint: disable=ungrouped-imports,no-name-in-module + flash_attn_func as flash_attn_func_v4, + flash_attn_varlen_func as flash_attn_varlen_func_v4, + _validate_head_dims as _fa4_validate_head_dims, + ) + except ImportError as exc: + flash_attn_func_v4 = None + flash_attn_varlen_func_v4 = None + warnings.warn( + f"FlashAttention 4 is installed but cannot be loaded: {exc}", + RuntimeWarning, + stacklevel=2, + ) + else: + fa_utils.v4_validate_head_dims = _fa4_validate_head_dims + fa_utils.set_flash_attention_4_params() # Float8CurrentScaling: fused_attn_bwd takes O in FP8 by default, this flag allows it in F16 _dpa_fp8_cs_o_in_f16 = os.getenv("NVTE_DPA_FP8CS_O_in_F16", "1") == "1" @@ -401,7 +422,22 @@ def fast_setattr(self, name: str, value: Any) -> None: """Fast attribute set for non-parameter fields.""" self.__dict__[name] = value - def forward( + def forward(self, *args, fp8: bool = False, fp8_output: bool = False, **kwargs) -> torch.Tensor: + """Unfused attention fprop; see `_forward` for the argument list. + + FP8 (emulation and/or Float8Tensor output) is not supported under + torch.compile -- run the backend as an eager island in that case. + """ + if fp8 or fp8_output: + return self._forward_eager(*args, fp8=fp8, fp8_output=fp8_output, **kwargs) + return self._forward(*args, fp8=False, fp8_output=False, **kwargs) + + @no_torch_dynamo() + def _forward_eager(self, *args, **kwargs) -> torch.Tensor: + """Eager-only (dynamo-disabled) wrapper around `_forward`.""" + return self._forward(*args, **kwargs) + + def _forward( self, _alibi_cache: Dict[str, Any], query_layer: torch.Tensor, @@ -436,6 +472,14 @@ def forward( if inference_params is not None and inference_params.is_paged: key_layer, value_layer = inference_params.convert_paged_to_nonpaged(self.layer_number) + # Token count for the thd output conversion (ConvertBSHDtoTHD) below. + # Captured here, before any layout conversion, because when the query + # enters in thd layout (qkv_format "thd" for training or "thd_2bshd" for + # inference) shape[0] is the total query token count. Deriving it later + # via cu_seqlens_q[-1].item() would sync with the GPU and break + # torch.compile + cudagraphs (unbacked SymInt). + total_tokens_q = query_layer.shape[0] if q_format == "thd" else None + # convert to sbhd # training: bshd, thd # inference: bshd, sbhd_2bshd, thd_2bshd @@ -744,6 +788,7 @@ def forward( context_layer = ConvertBSHDtoTHD.apply( context_layer, cu_seqlens_q, + total_tokens_q, ) # [tq, h, d] --> [tq, hd] @@ -825,8 +870,8 @@ def __init__( fa_utils.version >= fa_utils.version_required ), f"FlashAttention minimum version {fa_utils.version_required} is required." assert ( - fa_utils.version <= fa_utils.max_version - ), f"FlashAttention maximum version {fa_utils.max_version} is supported." + fa_utils.version < fa_utils.max_version + ), f"FlashAttention versions before {fa_utils.max_version} are supported." self.softmax_scale = softmax_scale self.attention_dropout_ctx = attention_dropout_ctx @@ -1051,6 +1096,18 @@ def forward( use_flash_attn_3 = ( flash_attention_backend is not None and flash_attention_backend.major == 3 ) + if ( + use_flash_attn_4 + and (10, 0) <= get_device_compute_capability() < (12, 0) + and query_layer.shape[-1] == key_layer.shape[-1] == value_layer.shape[-1] == 256 + and all(not isinstance(x, Float8Tensor) for x in [query_layer, key_layer, value_layer]) + and any(not x.is_contiguous() for x in [query_layer, key_layer, value_layer]) + ): + # FA4 D=256 SM10x kernels need packed K/V views materialized, even + # when the last dimension is contiguous. + query_layer, key_layer, value_layer = [ + x.contiguous() for x in (query_layer, key_layer, value_layer) + ] if context_parallel and all( not isinstance(x, Float8Tensor) for x in [query_layer, key_layer, value_layer] ): @@ -1362,8 +1419,11 @@ def forward( deterministic, softmax_offset, fp8_output, + bf16_backward, layer_number, return_max_logit, + packed_qkv=None, + packed_kv=None, ): # pylint: disable=missing-function-docstring @@ -1424,9 +1484,19 @@ def forward( # fp8_dtype = tex.DType.kFloat8E4M3 if is_input_fp8: q_fp8, k_fp8, v_fp8 = q, k, v + + if fp8_recipe.mxfp8(): + qkv_scale_inv_format = "bhsd" # Same as what combine_and_quantize would give else: q_fp8, k_fp8, v_fp8, qkv_layout, qkv_scale_inv_format = combine_and_quantize( - qkv_layout, q, k, v, QKV_quantizer, used_in_backward=is_training + qkv_layout, + q, + k, + v, + QKV_quantizer, + used_in_backward=is_training, + combined_qkv=packed_qkv, + combined_kv=packed_kv, ) # print quantizers @@ -1592,6 +1662,8 @@ def forward( ctx.is_input_fp8 = is_input_fp8 ctx.is_output_fp8 = is_output_fp8 + # Return dQ/dK/dV in bf16 even if is_input_fp8 + ctx.bf16_backward = bf16_backward tensors_to_save, tensor_objects = prepare_for_saving( *fp8_tensors, @@ -1858,7 +1930,8 @@ def backward(ctx, d_out, *_args): # dq, dk, dv: torch.Tensor; dtype = torch.float16 or torch.bfloat16 dq, dk, dv = dq_, dk_, dv_ is_quantized_tensor = isinstance(dq_, QuantizedTensorStorage) - if is_quantized_tensor and not ctx.is_input_fp8: + + if is_quantized_tensor and (not ctx.is_input_fp8 or ctx.bf16_backward): # return in F16 dq, dk, dv = combine_and_dequantize( ctx.dqkv_layout, @@ -1867,7 +1940,7 @@ def backward(ctx, d_out, *_args): dv_, src_nominal_dtype=dq_.dtype, ) - if not is_quantized_tensor and ctx.is_input_fp8: + if not is_quantized_tensor and ctx.is_input_fp8 and not ctx.bf16_backward: # return in FP8 dq, dk, dv, _, _ = combine_and_quantize( ctx.dqkv_layout, dq_, dk_, dv_, ctx.dQKV_quantizer @@ -1964,6 +2037,9 @@ def backward(ctx, d_out, *_args): None, None, None, + None, # packed_qkv + None, # packed_kv + None, ) @@ -2041,7 +2117,7 @@ def forward( attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None, window_size: Optional[Tuple[int, int]] = None, bottom_right_diagonal: Optional[bool] = None, - fused_attention_backend: tex.NVTE_Fused_Attn_Backend = tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend, + fused_attention_backend: FusedAttnBackend = FusedAttnBackend["No_Backend"], core_attention_bias_type: str = "no_bias", core_attention_bias: Optional[torch.Tensor] = None, fast_zero_fill: bool = True, @@ -2060,10 +2136,13 @@ def forward( score_mod_bprop: Optional[Callable] = None, score_mod_tensors: Optional[Dict[str, torch.Tensor]] = None, score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]] = None, + packed_qkv: Optional[torch.Tensor] = None, + packed_kv: Optional[torch.Tensor] = None, + bf16_backward: bool = False, ) -> torch.Tensor: """fused attention fprop""" assert ( - fused_attention_backend != tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend + fused_attention_backend != FusedAttnBackend["No_Backend"] ), "No fused attention backend supports this input combination!" assert all( x.dtype in [torch.float16, torch.bfloat16] or isinstance(x, QuantizedTensorStorage) @@ -2156,15 +2235,15 @@ def forward( use_FAv2_bwd = ( self.use_FAv2_bwd and (core_attention_bias_type == "no_bias") - and (fused_attention_backend == tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen) + and (fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"]) ) if fp8: fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8_meta is not None and fp8_meta.get("local_recipes", None) is not None: fp8_recipe = fp8_meta["local_recipes"][0] - assert fused_attention_backend == tex.NVTE_Fused_Attn_Backend.NVTE_FP8, ( - f"cuDNN attention sub-backend {int(tex.NVTE_Fused_Attn_Backend.NVTE_FP8)}" + assert fused_attention_backend == FusedAttnBackend["FP8"], ( + f"cuDNN attention sub-backend {int(FusedAttnBackend['FP8'])}" " is required for FP8 attention!" ) assert fp8_meta is not None, "FP8 metadata fp8_meta is required for FP8 attention!" @@ -2184,7 +2263,10 @@ def forward( if context_parallel: assert ( - IS_HIP_EXTENSION or fp8 + # ROCm CK/AOTriton backends support CP; NVTE_F16_arbitrary_seqlen is + # NVIDIA-only, so short-circuit on ROCm before touching that enumerator. + IS_HIP_EXTENSION + or fp8 or fused_attention_backend == tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen ), f"{fused_attention_backend} does not work with context parallelism!" assert core_attention_bias_type not in [ @@ -2277,8 +2359,11 @@ def forward( self.deterministic, softmax_offset, fp8_output, + bf16_backward, self.layer_number, self.return_max_logit, + packed_qkv, + packed_kv, ) if self.return_max_logit: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 53eea5e329..cbbf643b78 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -96,6 +96,19 @@ def get_bsh_dims(tensor_format): return batch_dim, seq_dim, head_dim +def _zero_thd_padding(tensors, cu_seqlens, cu_seqlens_padded): + """Zero inter-sequence padding in one or more tensors.""" + if cu_seqlens is None or cu_seqlens_padded is None: + return + tensor = next((tensor for tensor in tensors if tensor is not None), None) + if tensor is None: + return + padding_mask = dpa_utils.get_thd_padding_mask(tensor.shape[0], cu_seqlens, cu_seqlens_padded) + for tensor in tensors: + if tensor is not None: + tensor[padding_mask] = 0 + + def flash_attn_p2p_communicate( rank, send_tensor, send_dst, recv_tensor, recv_src, cp_group, batch_p2p_comm ): @@ -2663,8 +2676,8 @@ def backward(ctx, dout, *_args): elif ctx.qkv_format == "sbhd": dq[0].fill_(0) dq[1].copy_(dq_) - else: - dq.copy_(dq_) + elif ctx.qkv_format == "thd": + tex.thd_grad_correction(dq, dq_, cu_seqlens_q_padded, "zero", "copy") elif causal: if i > (cp_size - rank - 1): dq.add_(dq_) @@ -2750,9 +2763,9 @@ def backward(ctx, dout, *_args): dk[1].fill_(0) dv[0].copy_(dv_) dv[1].fill_(0) - else: - dk.copy_(dk_) - dv.copy_(dv_) + elif ctx.qkv_format == "thd": + tex.thd_grad_correction(dk, dk_, cu_seqlens_kv_padded, "copy", "zero") + tex.thd_grad_correction(dv, dv_, cu_seqlens_kv_padded, "copy", "zero") else: dk.copy_(dk_) dv.copy_(dv_) @@ -2889,6 +2902,12 @@ def backward(ctx, dout, *_args): ctx.dP_quantizer, ) + # Partial-gradient reduction can write THD inter-sequence padding. + # Clean it while gradients and per-step sequence metadata share sequence order. + if ctx.qkv_format == "thd": + _zero_thd_padding((dq,), cu_seqlens_q_per_step[0], cu_seqlens_q_padded) + _zero_thd_padding((dk, dv), cu_seqlens_kv_per_step[0], cu_seqlens_kv_padded) + if cp_size_a2a > 1: if ctx.fp8 and ctx.is_input_fp8: dq_fp8, dk_fp8, dv_fp8 = dq, dk, dv @@ -2926,23 +2945,6 @@ def backward(ctx, dout, *_args): nvtx_range_pop(f"{nvtx_label}") - # Zero-fill dQ/dK/dV at positions beyond the actual sequence end (THD CUDA Graph). - # cu_seqlens_*_padded are already local to this CP rank in the THD path. - # Use Q's padded boundary for dQ and KV's padded boundary for dK/dV. - # Skip the corresponding zero-fill when its padded cu_seqlens is absent. - if ctx.qkv_format == "thd": - if cu_seqlens_q_padded is not None and isinstance(dq, torch.Tensor) and dq.shape[0] > 0: - q_pad_mask = torch.arange(dq.shape[0], device=dq.device) >= cu_seqlens_q_padded[-1] - dq[q_pad_mask] = 0 - if cu_seqlens_kv_padded is not None: - kv_actual_t = cu_seqlens_kv_padded[-1] - for d_tensor in [dk, dv]: - if isinstance(d_tensor, torch.Tensor) and d_tensor.shape[0] > 0: - kv_pad_mask = ( - torch.arange(d_tensor.shape[0], device=d_tensor.device) >= kv_actual_t - ) - d_tensor[kv_pad_mask] = 0 - return ( None, dq, @@ -3090,10 +3092,12 @@ def forward( window_size == (-1, 0) or window_size == (-1, -1) or use_fused_attention + or use_flash_attn_3 or fa_utils.v2_3_plus ), ( "cp_comm_type='all_gather' only supports SWA through FusedAttention or FlashAttention" - f" >= 2.3. Found {use_fused_attention=} and {fa_utils.v2_3_plus=}." + f" >= 2.3. Found {use_fused_attention=}, {use_flash_attn_3=}, " + f"and {fa_utils.v2_3_plus=}." ) assert q.shape[seq_dim_qkv] % 2 == 0 and k.shape[seq_dim_qkv] % 2 == 0, ( "cp_comm_type='all_gather' requires seq_len % 2 == 0 for Q, K, V. Found seq_len_q =" @@ -3183,7 +3187,7 @@ def forward( fp8_meta_kwargs = {} if fp8: assert use_fused_attention, "FP8 is only supported with FusedAttention backend!" - fused_attn_backend = tex.NVTE_Fused_Attn_Backend.NVTE_FP8 + fused_attn_backend = FusedAttnBackend["FP8"] if not is_input_fp8 and not fp8_recipe.mxfp8(): q_fp8, k_fp8, v_fp8, qkv_layout, _ = combine_and_quantize( qkv_layout, q, k, v, QKV_quantizer @@ -3242,7 +3246,7 @@ def forward( # is large enough to outlast cp_stream's launch (e.g. bucket128k @ cp=8). cp_stream.wait_stream(torch.cuda.current_stream()) - # THD all_gather only reaches this path for f16/bf16 attention today. + # Shapes before per-step slicing and FP8 metadata wrapping. # q: [b, 2, s//2, h, d] or [2, s//2, b, h, d] # k: [s, b, h, d] # v: [s, b, h, d] @@ -3410,6 +3414,11 @@ def forward( ) max_seqlen_kv_ = kv_range[1] cu_seqlens_kv_per_step[i] = thd_cu_seqlens_kv_per_step[i] + if fp8: + q_part, k_part, v_part = [ + Float8Tensor.make_like(x, data=y, dtype=fwd_nominal_dtype) + for x, y in zip([q_fp8, k_fp8, v_fp8], [q_part, k_part, v_part]) + ] if use_fused_attention: # Set per-step parameters for THD vs bshd/sbhd if qkv_format == "thd": @@ -3742,7 +3751,8 @@ def backward(ctx, dout, *_args): # v: [s, b, h, d] if ctx.fp8 and not ctx.fp8_recipe.mxfp8(): q, k, v = [x._data for x in [q_fp8, k_fp8, v_fp8]] - if not ctx.qkv_reshaped: + # BSHD/SBHD split the sequence into two chunks; THD stays token-major [t, h, d]. + if not ctx.qkv_reshaped and ctx.qkv_format != "thd": q = q.view( *q.shape[:seq_dim_qkv], 2, q.shape[seq_dim_qkv] // 2, *q.shape[(seq_dim_qkv + 1) :] ) @@ -3919,7 +3929,7 @@ def backward(ctx, dout, *_args): qkv_scale_inv_format = None do_scale_inv_format = None if ctx.fp8: - fused_attn_backend = tex.NVTE_Fused_Attn_Backend.NVTE_FP8 + fused_attn_backend = FusedAttnBackend["FP8"] fp8_meta_kwargs["s_quantizer"] = ctx.S_quantizer fp8_meta_kwargs["dp_quantizer"] = ctx.dP_quantizer fp8_meta_kwargs["dqkv_quantizer"] = ctx.dQKV_quantizer @@ -4242,10 +4252,11 @@ def forward( window_size == (-1, 0) or window_size == (-1, -1) or use_fused_attention + or use_flash_attn_3 or fa_utils.v2_3_plus ), ( "cp_comm_type='a2a' only supports SWA through FusedAttention or FlashAttention >= 2.3." - f" Found {use_fused_attention=} and {fa_utils.v2_3_plus=}." + f" Found {use_fused_attention=}, {use_flash_attn_3=}, and {fa_utils.v2_3_plus=}." ) assert q.shape[seq_dim_qkv] % 2 == 0 and k.shape[seq_dim_qkv] % 2 == 0, ( "cp_comm_type='a2a' requires seq_len % 2 == 0 for Q, K, V. Found seq_len_q =" @@ -5029,9 +5040,6 @@ def attn_forward_func_with_cp( assert ( isinstance(cp_group, list) and len(cp_group) == 2 ), "CP implementation a2a+p2p requires cp_group = [a2a_cp_group, p2p_cp_group]!" - assert ( - qkv_format != "thd" - ), f"{qkv_format} format is not supported with hierarchical CP implementation yet!" assert ( attn_bias_type == "no_bias" ), f"{attn_bias_type} bias type is not supported with hierarchical CP implementation yet!" diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 9037ef4184..a4344cca3f 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -34,6 +34,7 @@ Float8BlockScalingRecipeState, ) from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import Float8TensorStorage +from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.export import is_in_onnx_export_mode from transformer_engine.pytorch.constants import AttnMaskTypes, AttnTypes, dist_group_type, DType @@ -89,6 +90,75 @@ "_alibi_bias_require_update": False, } + +def _infer_custom_dpa_local_recipes( + fp8_recipe: Recipe, + fp8_meta: Dict[str, Any], + quantizers: Dict[str, Any], +) -> Optional[List[Recipe]]: + """Infer native-equivalent DPA recipe labels for CustomRecipe control-flow. + + CustomRecipe owns quantizer construction, but DPA backend selection and a + few fused-attention branches still dispatch on recipe predicates. Attach + local recipe labels that match the qfactory DPA quantizer family while + keeping the actual qfactory-created quantizers untouched. + """ + try: + qkv_quantizer = quantizers["scaling_fwd"][dpa_utils.META_QKV] + except (KeyError, IndexError, TypeError): + return None + + from transformer_engine.pytorch.tensor.float8_tensor import ( + Float8CurrentScalingQuantizer, + Float8Quantizer, + ) + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + if isinstance(qkv_quantizer, MXFP8Quantizer): + return [ + MXFP8BlockScaling( + fp8_format=fp8_recipe.fp8_format, + fp8_dpa=fp8_recipe.fp8_dpa, + fp8_mha=fp8_recipe.fp8_mha, + ) + ] + + def _delayed_scaling_recipe() -> Optional[DelayedScaling]: + fwd_state = fp8_meta.get("scaling_fwd") + ds_recipe = getattr(fwd_state, "_inner_delayed_scaling_recipe", None) + if ds_recipe is None: + return None + return DelayedScaling( + fp8_format=ds_recipe.fp8_format, + margin=ds_recipe.margin, + amax_history_len=ds_recipe.amax_history_len, + amax_compute_algo=ds_recipe.amax_compute_algo, + scaling_factor_compute_algo=ds_recipe.scaling_factor_compute_algo, + reduce_amax=ds_recipe.reduce_amax, + fp8_dpa=fp8_recipe.fp8_dpa, + fp8_mha=fp8_recipe.fp8_mha, + ) + + if isinstance(qkv_quantizer, Float8CurrentScalingQuantizer): + ds_recipe = _delayed_scaling_recipe() + if ds_recipe is not None: + return [ + Float8CurrentScaling( + fp8_format=fp8_recipe.fp8_format, + fp8_dpa=fp8_recipe.fp8_dpa, + fp8_mha=fp8_recipe.fp8_mha, + ), + ds_recipe, + ] + + if isinstance(qkv_quantizer, Float8Quantizer): + ds_recipe = _delayed_scaling_recipe() + if ds_recipe is not None: + return [ds_recipe] + + return None + + """ This feature is **experimental** and subject to change. @@ -200,6 +270,111 @@ def _trim_output(attn_out, num_attention_heads, padded_head_dim_v, orig_head_dim return attn_out[..., :orig_head_dim_v].reshape(*out_shape, -1) +def _unpack_packed_qkv( + qkv_layer: Optional[torch.Tensor], + kv_layer: Optional[torch.Tensor], + query_layer: Optional[torch.Tensor], + key_layer: Optional[torch.Tensor], + value_layer: Optional[torch.Tensor], + qkv_format: str, + qkv_interleave_dim: int, + inference_params: Optional[InferenceParams], +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[str]]: + """Resolve declarative packed inputs into q/k/v. + + Derives q/k/v as zero-copy views of the packed buffer (``qkv_layer`` or + ``kv_layer``) and constructs the exact layout string from the declaration. + The layout enum is truthful by construction, so no pointer-based detection + is needed downstream (this also covers thd and FP8 DPA). + + Returns ``(query_layer, key_layer, value_layer, declared_qkv_layout)``; + the layout is ``None`` when no packed input is given. + """ + if qkv_layer is None and kv_layer is None: + if query_layer is None or key_layer is None or value_layer is None: + raise ValueError( + "query_layer, key_layer and value_layer are required unless packed" + " inputs (qkv_layer or query_layer + kv_layer) are provided." + ) + return query_layer, key_layer, value_layer, None + + if qkv_layer is not None and kv_layer is not None: + raise ValueError("qkv_layer and kv_layer are mutually exclusive.") + if inference_params is not None: + raise ValueError( + "Packed inputs (qkv_layer/kv_layer) are not supported with KV caching" + " (inference_params); pass separate query/key/value tensors instead." + ) + if qkv_interleave_dim not in (-3, -2): + raise ValueError( + "qkv_interleave_dim must be -3 (e.g. bs3hd) or -2 (e.g. bsh3d), got" + f" {qkv_interleave_dim}." + ) + packed = qkv_layer if qkv_layer is not None else kv_layer + # The declared layout describes the packed buffer's memory, so it must have + # stride 1 in its last dimension (the check get_qkv_layout would otherwise + # perform on the derived q/k/v views). + if packed.stride(-1) != 1: + raise ValueError( + "The packed tensor (qkv_layer/kv_layer) must have stride 1 in its last" + f" dimension, got strides {tuple(packed.stride())}." + ) + + def _packed_layout(fmt: str, num: int) -> str: + # bshd + 3 @ -3 -> bs3hd; bshd + 3 @ -2 -> bsh3d; thd + 2 @ -2 -> th2d + pos = len(fmt) + qkv_interleave_dim + 1 + return fmt[:pos] + str(num) + fmt[pos:] + + if qkv_layer is not None: + if any(x is not None for x in (query_layer, key_layer, value_layer)): + raise ValueError( + "qkv_layer already packs Q, K and V: query_layer, key_layer and" + " value_layer must be None when qkv_layer is provided." + ) + expected_rank = 4 if qkv_format == "thd" else 5 + if qkv_layer.dim() != expected_rank: + raise ValueError( + f"qkv_layer must be a {expected_rank}D tensor for" + f" qkv_format={qkv_format!r}, got {qkv_layer.dim()}D." + ) + if qkv_layer.shape[qkv_interleave_dim] != 3: + raise ValueError( + f"qkv_layer must have size 3 at dim {qkv_interleave_dim}" + f" (qkv_interleave_dim), got shape {tuple(qkv_layer.shape)}." + ) + query_layer, key_layer, value_layer = ( + qkv_layer.select(qkv_interleave_dim, i) for i in range(3) + ) + return query_layer, key_layer, value_layer, _packed_layout(qkv_format, 3) + + if query_layer is None: + raise ValueError( + "kv_layer packs only K and V: query_layer is required when kv_layer is provided." + ) + if key_layer is not None or value_layer is not None: + raise ValueError( + "kv_layer already packs K and V: key_layer and value_layer must be" + " None when kv_layer is provided." + ) + if kv_layer.dim() != query_layer.dim() + 1: + raise ValueError( + "kv_layer must have one more dimension than query_layer, got" + f" {kv_layer.dim()}D kv_layer and {query_layer.dim()}D query_layer." + ) + if kv_layer.shape[qkv_interleave_dim] != 2: + raise ValueError( + f"kv_layer must have size 2 at dim {qkv_interleave_dim}" + f" (qkv_interleave_dim), got shape {tuple(kv_layer.shape)}." + ) + key_layer, value_layer = (kv_layer.select(qkv_interleave_dim, i) for i in range(2)) + return ( + query_layer, + key_layer, + value_layer, + f"{qkv_format}_{_packed_layout(qkv_format, 2)}", + ) + + class DotProductAttention(TransformerEngineBaseModule): r"""Allows the model to jointly attend to information from different representation subspaces as described in the paper: @@ -223,6 +398,90 @@ class DotProductAttention(TransformerEngineBaseModule): As the FP8 attention support expands from one backend to multiple backends, the location of that key has also shifted (see `FP8 checkpoint compatibility `_). + .. rubric:: Fine-grained Linear and attention recipes + + .. warning:: + + Fine-grained attention configuration through ``CustomRecipe`` and a quantizer factory is + experimental and subject to change. + + A quantizer factory can select different recipes for Linear and DotProductAttention tensors + using ``QuantizerRole``. DotProductAttention itself supports only its fixed recipe families: + FP8 delayed scaling, FP8 current scaling, and MXFP8 block scaling. In particular, a factory may + return NVFP4 quantizers for Linear roles, but it must not return NVFP4 quantizers for DPA roles. + + .. list-table:: Example Linear and attention combinations + :header-rows: 1 + + * - Linear + - Attention + - Configuration + - Status + * - NVFP4 + - FP8 current scaling for QKV/O and delayed scaling for S/dP + - ``nvfp4_linear_fp8_dpa_factory`` + - Validated factory provided by Transformer Engine + * - NVFP4 + - MXFP8 + - User-defined factory shown below + - Experimental example; not broadly validated + + The validated NVFP4 Linear + FP8 attention combination is available from the quantizer factory + zoo:: + + from transformer_engine.common.recipe import CustomRecipe + from transformer_engine.pytorch.quantization import autocast + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + nvfp4_linear_fp8_dpa_factory, + ) + + recipe = CustomRecipe( + qfactory=nvfp4_linear_fp8_dpa_factory, + fp8_dpa=True, + ) + with autocast(recipe=recipe): + output = model(input) + + The following factory demonstrates the experimental NVFP4 Linear + MXFP8 attention + combination. With ``CustomRecipe``, the per-role selection is expressed directly in the + factory, so ``NVTE_DPA_FP8_RECIPE`` is not needed. DPA also issues hint-only roles for its + output boundaries; these must resolve to a DPA-supported quantizer even when the boundary + tensor remains in BF16. For MXFP8 attention, the fused kernel handles the S/dP slots + internally, so their factory-provided quantizers are not consumed:: + + from transformer_engine.common.recipe import CustomRecipe + from transformer_engine.pytorch.constants import DType + from transformer_engine.pytorch.custom_recipes.quantizer_factories import nvfp4_factory + from transformer_engine.pytorch.quantization import autocast + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + def nvfp4_linear_mxfp8_dpa_factory(role): + # NVFP4 for Linear roles and MXFP8 for supported DPA roles. + is_dpa = role is not None and role.module_type == "dpa" + is_dpa_boundary = ( + role is not None + and not role.module_type + and ("dpa_output" in role.name or "dpa_grad_input" in role.name) + ) + + if is_dpa or is_dpa_boundary: + is_bwd_role = ( + is_dpa and role.tensor_type in ("do", "dp", "dqkv") + ) or ( + is_dpa_boundary and "dpa_grad_input" in role.name + ) + fp8_dtype = DType.kFloat8E5M2 if is_bwd_role else DType.kFloat8E4M3 + return MXFP8Quantizer(fp8_dtype=fp8_dtype) + + return nvfp4_factory(role) + + recipe = CustomRecipe( + qfactory=nvfp4_linear_mxfp8_dpa_factory, + fp8_dpa=True, + ) + with autocast(recipe=recipe): + output = model(input) + Parameters ---------- @@ -400,6 +659,14 @@ def __init__( assert not return_max_logit, "ROCm does not support return_max_logit yet." super().__init__(name=name) + # Cache the native recipe labels inferred from custom DPA quantizers. + # ``init_fp8_metadata`` runs on every forward, while the quantizers only + # change when their recipe state is rebuilt. + self._custom_dpa_local_recipes_cache_key: Optional[Tuple[Any, ...]] = None + self._custom_dpa_local_recipes_cache: Optional[List[Recipe]] = None + self._qkv_capabilities_quantizer: Optional[Any] = None + self._qkv_capabilities_cache: Optional[Tuple[bool, bool]] = None + self.logger = logging.getLogger("DotProductAttention") self.logger.setLevel(attn_log._log_level) if not self.logger.hasHandlers(): @@ -620,6 +887,26 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8_recipe.custom(): super().init_fp8_metadata(num_gemms=num_gemms) + fwd_quantizers = self.quantizers.get("scaling_fwd", ()) + cache_key = ( + id(self.fp8_meta.get("scaling_fwd")), + tuple(id(quantizer) for quantizer in fwd_quantizers), + fp8_recipe.fp8_format, + fp8_recipe.fp8_dpa, + fp8_recipe.fp8_mha, + ) + if cache_key != self._custom_dpa_local_recipes_cache_key: + self._custom_dpa_local_recipes_cache = _infer_custom_dpa_local_recipes( + fp8_recipe, self.fp8_meta, self.quantizers + ) + self._custom_dpa_local_recipes_cache_key = cache_key + + if self._custom_dpa_local_recipes_cache is None: + # Do not leave labels from an earlier supported quantizer + # family attached after a rebuild to an unsupported family. + self.fp8_meta.pop("local_recipes", None) + else: + self.fp8_meta["local_recipes"] = self._custom_dpa_local_recipes_cache return # switch/append recipe: fp8_recipe stays unchanged, but DPA.fp8_meta["recipe"] may be set to @@ -826,6 +1113,57 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: # Clear cached workspaces as they were created with the old recipe/quantizer type self._fp8_workspaces.clear() + def get_qkv_quantization_capabilities(self) -> Tuple[bool, bool]: + """Return MHA boundary capabilities from the canonical QKV quantizer. + + The returned flags are ``(float8_current_scaling, mxfp8_scaling)``. + """ + self.init_fp8_metadata(num_gemms=3) + try: + qkv_quantizer = self.quantizers["scaling_fwd"][dpa_utils.META_QKV] + except (KeyError, IndexError, TypeError) as exc: + role = QuantizerRole( + module_type="dpa", + tensor_type="qkv", + name=self.name or "", + ) + raise RuntimeError( + f"DotProductAttention did not materialize the canonical QKV quantizer for {role}." + ) from exc + + if qkv_quantizer is self._qkv_capabilities_quantizer: + assert self._qkv_capabilities_cache is not None + return self._qkv_capabilities_cache + + from transformer_engine.pytorch.tensor.float8_tensor import ( + Float8CurrentScalingQuantizer, + Float8Quantizer, + ) + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + if isinstance(qkv_quantizer, Float8CurrentScalingQuantizer): + capabilities = (True, False) + elif isinstance(qkv_quantizer, MXFP8Quantizer): + capabilities = (False, True) + elif isinstance(qkv_quantizer, Float8Quantizer): + capabilities = (False, False) + else: + capabilities = None + + if capabilities is not None: + self._qkv_capabilities_quantizer = qkv_quantizer + self._qkv_capabilities_cache = capabilities + return capabilities + + role = QuantizerRole( + module_type="dpa", + tensor_type="qkv", + name=self.name or "", + ) + raise TypeError( + f"Unsupported CustomRecipe quantizer for {role}: {type(qkv_quantizer).__name__}." + ) + def set_meta_tensor(self, fwd: bool, recipe: Union[Recipe, List[Recipe]]) -> None: """Override to allow multiple recipes. Init scales and amaxes for fwd | bwd.""" if isinstance(recipe, Recipe) and recipe.custom(): @@ -994,9 +1332,9 @@ def get_quantizer_roles( @no_torch_dynamo(recursive=False) def forward( self, - query_layer: torch.Tensor, - key_layer: torch.Tensor, - value_layer: torch.Tensor, + query_layer: Optional[torch.Tensor] = None, + key_layer: Optional[torch.Tensor] = None, + value_layer: Optional[torch.Tensor] = None, attention_mask: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]] = None, qkv_format: str = None, cu_seqlens_q: torch.Tensor = None, @@ -1016,11 +1354,15 @@ def forward( inference_params: Optional[InferenceParams] = None, pad_between_seqs: Optional[bool] = None, fp8_output: Optional[bool] = False, + bf16_backward: Optional[bool] = False, num_splits: Optional[int] = 1, score_mod: Optional[Callable] = None, score_mod_bprop: Optional[Callable] = None, score_mod_tensors: Optional[Dict[str, torch.Tensor]] = None, score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]] = None, + qkv_layer: Optional[torch.Tensor] = None, + kv_layer: Optional[torch.Tensor] = None, + qkv_interleave_dim: int = -3, ) -> torch.Tensor: r""" Dot Product Attention Layer. @@ -1119,12 +1461,15 @@ def forward( Parameters ---------- - query_layer : torch.Tensor - Query tensor. - key_layer : torch.Tensor - Key tensor. - value_layer : torch.Tensor - Value tensor. + query_layer : Optional[torch.Tensor], default = None + Query tensor. Required unless a packed input (``qkv_layer``, or + ``kv_layer`` together with ``query_layer``) is provided instead. + key_layer : Optional[torch.Tensor], default = None + Key tensor. Required unless a packed input (``qkv_layer`` or ``kv_layer``) + is provided instead. + value_layer : Optional[torch.Tensor], default = None + Value tensor. Required unless a packed input (``qkv_layer`` or ``kv_layer``) + is provided instead. attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]], default = None. Boolean tensor(s) used to mask out attention softmax input. It should be ``None`` for causal masks and ``"no_mask"``. For padding masks, it should be @@ -1232,8 +1577,44 @@ def forward( Runtime tensors exposed to score_mod_bprop as cuDNN graph tensors. Keys are user-defined string names consumed by the callback through ``tensors[name]``; there is no predefined set of accepted keys. + qkv_layer: Optional[torch.Tensor], default = None + Fully packed QKV tensor. When the QKV projection produces one packed buffer + (e.g. a fused QKV GEMM), it can be passed here directly instead of slicing + it into :attr:`query_layer`/:attr:`key_layer`/:attr:`value_layer` views. + For :attr:`qkv_format` = {"bshd", "sbhd"}, it must be a 5D tensor of shape + ``[b, s, 3, h, d]``/``[s, b, 3, h, d]`` (:attr:`qkv_interleave_dim` = -3) or + ``[b, s, h, 3, d]``/``[s, b, h, 3, d]`` (:attr:`qkv_interleave_dim` = -2); + for :attr:`qkv_format` = "thd", a 4D tensor of shape ``[t, 3, h, d]`` or + ``[t, h, 3, d]``. Q/K/V are derived as zero-copy views and the memory layout + (e.g. ``bs3hd``) is declared from the packing itself, so no pointer-based + layout detection runs on this path -- including for "thd" and FP8 attention. + Mutually exclusive with :attr:`query_layer`, :attr:`key_layer`, + :attr:`value_layer` and :attr:`kv_layer`. + kv_layer: Optional[torch.Tensor], default = None + Packed KV tensor, used together with :attr:`query_layer` + (e.g. ``[b, s, 2, hg, d]`` for :attr:`qkv_interleave_dim` = -3, or + ``[b, s, hg, 2, d]`` for :attr:`qkv_interleave_dim` = -2). K/V are derived + as zero-copy views and the layout (e.g. ``bshd_bs2hd``) is declared, not + detected. Mutually exclusive with :attr:`key_layer`, :attr:`value_layer` + and :attr:`qkv_layer`. + qkv_interleave_dim: int, default = -3 + Dimension of :attr:`qkv_layer`/:attr:`kv_layer` where the 3 (QKV) or 2 (KV) + interleave sits; must be -3 (e.g. ``bs3hd``) or -2 (e.g. ``bsh3d``, + Megatron-style). This is an explicit knob rather than shape inference, + since e.g. ``h == 3`` would make the shapes ambiguous. """ + query_layer, key_layer, value_layer, declared_qkv_layout = _unpack_packed_qkv( + qkv_layer, + kv_layer, + query_layer, + key_layer, + value_layer, + qkv_format if qkv_format is not None else self.qkv_format, + qkv_interleave_dim, + inference_params, + ) + with self.prepare_forward_ctx( query_layer, num_gemms=3, @@ -1419,7 +1800,15 @@ def forward( cu_seqlens_kv_padded = None # get qkv's memory layout - if all( + if declared_qkv_layout is not None: + # Packed inputs (qkv_layer/kv_layer) declare the layout: the enum is + # truthful by construction, so the pointer-based detection in + # get_qkv_layout is skipped entirely -- for dense, thd (t3hd/th3d) + # and FP8 DPA alike. + qkv_layout = declared_qkv_layout + q_format = qkv_format + kv_format = qkv_format + elif all( isinstance(x, Float8TensorStorage) for x in [query_layer, key_layer, value_layer] ): ( @@ -1436,6 +1825,25 @@ def forward( qkv_format=qkv_format, inference_params=inference_params, ) + elif all( + isinstance(x, MXFP8TensorStorage) for x in [query_layer, key_layer, value_layer] + ): + # Pre-quantized MXFP8 q/k/v: the wrapper has no real storage, so run + # layout detection on the underlying rowwise data (mirrors the Float8 path). + ( + qkv_layout, + query_layer._rowwise_data, + key_layer._rowwise_data, + value_layer._rowwise_data, + q_format, + kv_format, + ) = dpa_utils.get_qkv_layout( + query_layer._rowwise_data, + key_layer._rowwise_data, + value_layer._rowwise_data, + qkv_format=qkv_format, + inference_params=inference_params, + ) else: ( qkv_layout, @@ -1807,6 +2215,9 @@ def forward( inference_params=inference_params, softmax_offset=softmax_offset, fp8_output=fp8_output, + packed_qkv=qkv_layer, + packed_kv=kv_layer, + bf16_backward=bf16_backward, ) return self.fused_attention( query_layer, @@ -1842,6 +2253,9 @@ def forward( score_mod_bprop=score_mod_bprop, score_mod_tensors=score_mod_tensors, score_mod_bprop_tensors=score_mod_bprop_tensors, + packed_qkv=qkv_layer, + packed_kv=kv_layer, + bf16_backward=bf16_backward, ) if use_unfused_attention: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/softmax.py b/transformer_engine/pytorch/attention/dot_product_attention/softmax.py index 74d9583ce5..6e8a14402f 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/softmax.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/softmax.py @@ -3,142 +3,248 @@ # See LICENSE for license information. """Fused scaled masked softmax functions""" + import os -from typing import Callable, Tuple, Union, Optional +from typing import Callable, Optional import torch from torch import nn import transformer_engine_torch as tex from transformer_engine.pytorch.export import is_in_onnx_export_mode - THREADS_PER_WARP = 32 THREADS_PER_BLOCK = 128 -_default_causal_mask = {} +# ----------------------------- ScaledSoftmax ------------------------------- -def _get_default_causal_mask(mask_type: str, sq: int, sk: int) -> torch.Tensor: - """Return the causal upper triangular mask for softmax input""" +@torch.library.custom_op("te_softmax::scaled_softmax_fwd", mutates_args=()) +def scaled_softmax_forward(inputs: torch.Tensor, scale: float) -> torch.Tensor: + """Forward pass for ScaledSoftmax.""" + return tex.scaled_softmax_forward(inputs, scale) - def _get_mask(): - diagonal_offset = sk - sq + 1 if "bottom_right" in mask_type else 1 - return torch.triu( - torch.ones(sq, sk, dtype=torch.bool, device="cuda"), diagonal=diagonal_offset - ) - if is_in_onnx_export_mode(): - return _get_mask() - matrix_identifiers = (mask_type, sq, sk) - if matrix_identifiers not in _default_causal_mask: - _default_causal_mask[matrix_identifiers] = _get_mask() - return _default_causal_mask[matrix_identifiers] +@scaled_softmax_forward.register_fake +def _scaled_softmax_forward_fake(inputs: torch.Tensor, scale: float) -> torch.Tensor: + del scale + return torch.empty_like(inputs) -class ScaledUpperTriangMaskedSoftmax(torch.autograd.Function): - """ - Fused operation which performs following three operations in sequence - 1. Scale the tensor. - 2. Apply upper triangular mask (typically used in gpt models). - 3. Perform softmax. - """ +@torch.library.custom_op("te_softmax::scaled_softmax_bwd", mutates_args=()) +def scaled_softmax_backward( + output_grads: torch.Tensor, softmax_results: torch.Tensor, scale: float +) -> torch.Tensor: + """Backward pass for ScaledSoftmax.""" + return tex.scaled_softmax_backward(output_grads, softmax_results, scale) - @staticmethod - def forward(ctx, inputs: torch.Tensor, scale: float) -> torch.Tensor: - """ScaledUpperTriangMaskedSoftmax fwd""" - scale_t = torch.tensor([scale]) - softmax_results = tex.scaled_upper_triang_masked_softmax_forward(inputs, scale_t[0]) - ctx.save_for_backward(softmax_results, scale_t) - return softmax_results +@scaled_softmax_backward.register_fake +def _scaled_softmax_backward_fake( + output_grads: torch.Tensor, softmax_results: torch.Tensor, scale: float +) -> torch.Tensor: + del softmax_results, scale + return torch.empty_like(output_grads) - @staticmethod - def backward(ctx, output_grads: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]: - """ScaledUpperTriangMaskedSoftmax bwd""" - softmax_results, scale_t = ctx.saved_tensors - input_grads = tex.scaled_upper_triang_masked_softmax_backward( - output_grads, softmax_results, scale_t[0] - ) - return input_grads, None +def _scaled_softmax_setup_context(ctx, inputs, output): + _inp, scale = inputs + ctx.scale = scale + ctx.save_for_backward(output) -class ScaledAlignedCausalMaskedSoftmax(torch.autograd.Function): - """ - Fused operation which performs following three operations in sequence - 1. Scale the tensor. - 2. Apply causal mask aligned to the bottom right corner of the input matrix - 3. Perform softmax. - """ +def _scaled_softmax_backward_wrapper(ctx, grad_output): + (softmax_results,) = ctx.saved_tensors + grad_inputs = torch.ops.te_softmax.scaled_softmax_bwd(grad_output, softmax_results, ctx.scale) + return grad_inputs, None - @staticmethod - def forward(ctx, inputs: torch.Tensor, scale: float) -> torch.Tensor: - """ScaledAlignedCausalMaskedSoftmax fwd""" - scale_t = torch.tensor([scale]) - softmax_results = tex.scaled_aligned_causal_masked_softmax_forward(inputs, scale_t[0]) - ctx.save_for_backward(softmax_results, scale_t) - return softmax_results - @staticmethod - def backward(ctx, output_grads: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]: - """ScaledAlignedCausalMaskedSoftmax bwd""" - softmax_results, scale_t = ctx.saved_tensors - input_grads = tex.scaled_aligned_causal_masked_softmax_backward( - output_grads, softmax_results, scale_t[0] - ) +scaled_softmax_forward.register_autograd( + _scaled_softmax_backward_wrapper, + setup_context=_scaled_softmax_setup_context, +) - return input_grads, None +# --------------------------- ScaledMaskedSoftmax --------------------------- -class ScaledMaskedSoftmax(torch.autograd.Function): - """ - Fused operation which performs following three operations in sequence - 1. Scale the tensor. - 2. Apply the mask. - 3. Perform softmax. - """ - @staticmethod - def forward(ctx, inputs: torch.Tensor, mask: torch.Tensor, scale: float) -> torch.Tensor: - """ScaledMaskedSoftmax fwd""" - scale_t = torch.tensor([scale]) +@torch.library.custom_op("te_softmax::scaled_masked_softmax_fwd", mutates_args=()) +def scaled_masked_softmax_forward( + inputs: torch.Tensor, mask: torch.Tensor, scale: float +) -> torch.Tensor: + """Forward pass for ScaledMaskedSoftmax.""" + return tex.scaled_masked_softmax_forward(inputs, mask, scale) - softmax_results = tex.scaled_masked_softmax_forward(inputs, mask, scale_t[0]) - ctx.save_for_backward(softmax_results, scale_t) - return softmax_results - @staticmethod - def backward(ctx, output_grads: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]: - """ScaledMaskedSoftmax bwd""" - softmax_results, scale_t = ctx.saved_tensors +@scaled_masked_softmax_forward.register_fake +def _scaled_masked_softmax_forward_fake( + inputs: torch.Tensor, mask: torch.Tensor, scale: float +) -> torch.Tensor: + del mask, scale + return torch.empty_like(inputs) - input_grads = tex.scaled_masked_softmax_backward(output_grads, softmax_results, scale_t[0]) - return input_grads, None, None +@torch.library.custom_op("te_softmax::scaled_masked_softmax_bwd", mutates_args=()) +def scaled_masked_softmax_backward( + output_grads: torch.Tensor, softmax_results: torch.Tensor, scale: float +) -> torch.Tensor: + """Backward pass for ScaledMaskedSoftmax.""" + return tex.scaled_masked_softmax_backward(output_grads, softmax_results, scale) -class ScaledSoftmax(torch.autograd.Function): - """ - Fused operation which performs following two operations in sequence - 1. Scale the tensor. - 2. Perform softmax. - """ - @staticmethod - def forward(ctx, inputs: torch.Tensor, scale: float) -> torch.Tensor: - """ScaledSoftmax fwd""" - scale_t = torch.tensor([scale]) +@scaled_masked_softmax_backward.register_fake +def _scaled_masked_softmax_backward_fake( + output_grads: torch.Tensor, softmax_results: torch.Tensor, scale: float +) -> torch.Tensor: + del softmax_results, scale + return torch.empty_like(output_grads) - softmax_results = tex.scaled_softmax_forward(inputs, scale_t[0]) - ctx.save_for_backward(softmax_results, scale_t) - return softmax_results - @staticmethod - def backward(ctx, output_grads: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]: - """ScaledSoftmax bwd""" - softmax_results, scale_t = ctx.saved_tensors +def _scaled_masked_softmax_setup_context(ctx, inputs, output): + _inp, _mask, scale = inputs + ctx.scale = scale + ctx.save_for_backward(output) + + +def _scaled_masked_softmax_backward_wrapper(ctx, grad_output): + (softmax_results,) = ctx.saved_tensors + grad_inputs = torch.ops.te_softmax.scaled_masked_softmax_bwd( + grad_output, softmax_results, ctx.scale + ) + return grad_inputs, None, None + + +scaled_masked_softmax_forward.register_autograd( + _scaled_masked_softmax_backward_wrapper, + setup_context=_scaled_masked_softmax_setup_context, +) + + +# ---------------------- ScaledUpperTriangMaskedSoftmax ---------------------- + + +@torch.library.custom_op("te_softmax::scaled_upper_triang_masked_softmax_fwd", mutates_args=()) +def scaled_upper_triang_masked_softmax_forward(inputs: torch.Tensor, scale: float) -> torch.Tensor: + """Forward pass for ScaledUpperTriangMaskedSoftmax.""" + return tex.scaled_upper_triang_masked_softmax_forward(inputs, scale) + + +@scaled_upper_triang_masked_softmax_forward.register_fake +def _scaled_upper_triang_masked_softmax_forward_fake( + inputs: torch.Tensor, scale: float +) -> torch.Tensor: + del scale + return torch.empty_like(inputs) + - input_grads = tex.scaled_softmax_backward(output_grads, softmax_results, scale_t[0]) - return input_grads, None, None +@torch.library.custom_op("te_softmax::scaled_upper_triang_masked_softmax_bwd", mutates_args=()) +def scaled_upper_triang_masked_softmax_backward( + output_grads: torch.Tensor, softmax_results: torch.Tensor, scale: float +) -> torch.Tensor: + """Backward pass for ScaledUpperTriangMaskedSoftmax.""" + return tex.scaled_upper_triang_masked_softmax_backward(output_grads, softmax_results, scale) + + +@scaled_upper_triang_masked_softmax_backward.register_fake +def _scaled_upper_triang_masked_softmax_backward_fake( + output_grads: torch.Tensor, softmax_results: torch.Tensor, scale: float +) -> torch.Tensor: + del softmax_results, scale + return torch.empty_like(output_grads) + + +def _scaled_upper_triang_masked_softmax_setup_context(ctx, inputs, output): + _inp, scale = inputs + ctx.scale = scale + ctx.save_for_backward(output) + + +def _scaled_upper_triang_masked_softmax_backward_wrapper(ctx, grad_output): + (softmax_results,) = ctx.saved_tensors + grad_inputs = torch.ops.te_softmax.scaled_upper_triang_masked_softmax_bwd( + grad_output, softmax_results, ctx.scale + ) + return grad_inputs, None + + +scaled_upper_triang_masked_softmax_forward.register_autograd( + _scaled_upper_triang_masked_softmax_backward_wrapper, + setup_context=_scaled_upper_triang_masked_softmax_setup_context, +) + + +# -------------------- ScaledAlignedCausalMaskedSoftmax --------------------- + + +@torch.library.custom_op("te_softmax::scaled_aligned_causal_masked_softmax_fwd", mutates_args=()) +def scaled_aligned_causal_masked_softmax_forward( + inputs: torch.Tensor, scale: float +) -> torch.Tensor: + """Forward pass for ScaledAlignedCausalMaskedSoftmax.""" + return tex.scaled_aligned_causal_masked_softmax_forward(inputs, scale) + + +@scaled_aligned_causal_masked_softmax_forward.register_fake +def _scaled_aligned_causal_masked_softmax_forward_fake( + inputs: torch.Tensor, scale: float +) -> torch.Tensor: + del scale + return torch.empty_like(inputs) + + +@torch.library.custom_op("te_softmax::scaled_aligned_causal_masked_softmax_bwd", mutates_args=()) +def scaled_aligned_causal_masked_softmax_backward( + output_grads: torch.Tensor, softmax_results: torch.Tensor, scale: float +) -> torch.Tensor: + """Backward pass for ScaledAlignedCausalMaskedSoftmax.""" + return tex.scaled_aligned_causal_masked_softmax_backward(output_grads, softmax_results, scale) + + +@scaled_aligned_causal_masked_softmax_backward.register_fake +def _scaled_aligned_causal_masked_softmax_backward_fake( + output_grads: torch.Tensor, softmax_results: torch.Tensor, scale: float +) -> torch.Tensor: + del softmax_results, scale + return torch.empty_like(output_grads) + + +def _scaled_aligned_causal_masked_softmax_setup_context(ctx, inputs, output): + _inp, scale = inputs + ctx.scale = scale + ctx.save_for_backward(output) + + +def _scaled_aligned_causal_masked_softmax_backward_wrapper(ctx, grad_output): + (softmax_results,) = ctx.saved_tensors + grad_inputs = torch.ops.te_softmax.scaled_aligned_causal_masked_softmax_bwd( + grad_output, softmax_results, ctx.scale + ) + return grad_inputs, None + + +scaled_aligned_causal_masked_softmax_forward.register_autograd( + _scaled_aligned_causal_masked_softmax_backward_wrapper, + setup_context=_scaled_aligned_causal_masked_softmax_setup_context, +) + + +_default_causal_mask = {} + + +def _get_default_causal_mask(mask_type: str, sq: int, sk: int) -> torch.Tensor: + """Return the causal upper triangular mask for softmax input""" + + def _get_mask(): + diagonal_offset = sk - sq + 1 if "bottom_right" in mask_type else 1 + return torch.triu( + torch.ones(sq, sk, dtype=torch.bool, device="cuda"), diagonal=diagonal_offset + ) + + if is_in_onnx_export_mode(): + return _get_mask() + matrix_identifiers = (mask_type, sq, sk) + if matrix_identifiers not in _default_causal_mask: + _default_causal_mask[matrix_identifiers] = _get_mask() + return _default_causal_mask[matrix_identifiers] class FusedScaleMaskSoftmax(nn.Module): @@ -234,16 +340,16 @@ def forward_fused_softmax( padding, padding_causal, padding_causal_bottom_right | ScaledMaskedSoftmax arbitrary ([1, 1, sq, sk] or [b, 1, sq, sk]) | ScaledMaskedSoftmax """ - scale = 1.0 if scale is None else scale + scale = 1.0 if scale is None else float(scale) # Disable for now until unalignment bug is fixed. # if self.attn_mask_type in ["causal", "causal_bottom_right"]: - # return ScaledAlignedCausalMaskedSoftmax.apply(inp, scale) + # return torch.ops.te_softmax.scaled_aligned_causal_masked_softmax_fwd(inp, scale) # input is 4D tensor (1, 1, sq, sk) or (b, 1, sq, sk) if mask is not None and self.attn_mask_type != "no_mask": - return ScaledMaskedSoftmax.apply(inp, mask, scale) - return ScaledSoftmax.apply(inp, scale) + return torch.ops.te_softmax.scaled_masked_softmax_fwd(inp, mask, scale) + return torch.ops.te_softmax.scaled_softmax_fwd(inp, scale) def forward_torch_softmax( self, inp: torch.Tensor, mask: torch.Tensor, scale: Optional[float] = None diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index ae509302c3..8319f11ad7 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -7,6 +7,7 @@ """ Utils/Helper classes and methods for attention """ + import math import os from typing import Any, Callable, Dict, List, Optional, Tuple, Union @@ -38,6 +39,7 @@ META_DP, ) from transformer_engine.pytorch.attention.inference import InferenceParams +from transformer_engine.pytorch.cpu_offload import is_cpu_offload_enabled from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage from transformer_engine.pytorch.tensor.float8_tensor import ( Float8Tensor, @@ -106,7 +108,7 @@ def _get_supported_versions(version_min, version_max): """ Calculate version info based on min and max numbers """ - return ">= " + str(version_min) + ", " + "<= " + str(version_max) + return ">= " + str(version_min) + ", " + "< " + str(version_max) def maybe_contiguous(tensor: torch.Tensor) -> torch.Tensor: @@ -123,7 +125,7 @@ class FlashAttentionUtils: version = PkgVersion("0") version_required = PkgVersion("2.1.1") version_required_blackwell = PkgVersion("2.7.3") - max_version = PkgVersion("2.8.3") + max_version = PkgVersion("2.8.4") v2_plus = False v2_1_plus = False v2_3_plus = False @@ -157,6 +159,11 @@ class FlashAttentionUtils: # which raises AssertionError for unsupported (head_dim, head_dim_v) combinations. v4_validate_head_dims: Callable = None + @staticmethod + def is_version_supported(version: PkgVersion, minimum_version: PkgVersion) -> bool: + """Check whether a Flash Attention v2 version is supported.""" + return minimum_version <= version < FlashAttentionUtils.max_version + @staticmethod def set_flash_attention_version(): """ @@ -328,6 +335,59 @@ def __eq__(self, other): return True +class _NoOpLogger: + """ + Stand-in for the "DotProductAttention" logger used when get_attention_backend + is traced by torch.compile. logging.Logger methods are not traceable by dynamo + (they cause graph breaks), while this class's no-op methods are inlined away. + """ + + def debug(self, *args, **kwargs): + """No-op.""" + + def info(self, *args, **kwargs): + """No-op.""" + + def warning(self, *args, **kwargs): + """No-op.""" + + def error(self, *args, **kwargs): + """No-op.""" + + +_no_op_logger = _NoOpLogger() + + +@torch.compiler.assume_constant_result +def _get_fused_attn_backend( + is_training, + q_type, + kv_type, + qkv_layout, + bias_type, + attn_mask_type, + softmax_type, + *args, +): + """Constant-foldable tex.get_fused_attn_backend: the result depends only on + the attention config, and the python-side enum keeps it traceable by + torch.compile (see the FusedAttnBackend docstring). Layout/bias/mask/softmax + are taken as their string keys and resolved to the pybind enums here, so + that every argument is a python literal or a python enum.""" + return FusedAttnBackend.cast( + tex.get_fused_attn_backend( + is_training, + q_type, + kv_type, + QKVLayout[qkv_layout], + AttnBiasType[bias_type], + AttnMaskType[attn_mask_type], + SoftmaxType[softmax_type], + *args, + ) + ) + + def get_attention_backend( attention_params: AttentionParams = None, ): @@ -346,7 +406,7 @@ def get_attention_backend( If `use_flash_attention = True`, the version of the selected `FlashAttention` backend. use_fused_attention : bool Whether the `FusedAttention` backend has been selected. - fused_attention_backend : tex.NVTE_Fused_Attn_Backend + fused_attention_backend : FusedAttnBackend If `use_fused_attention = True`, one of `FusedAttention` three sub-backends, else `None`. use_unfused_attention : bool Whether the `UnfusedDotProductAttention` backend has been selected. @@ -393,17 +453,29 @@ def get_attention_backend( has_score_mod = attention_params.has_score_mod has_score_mod_bprop = attention_params.has_score_mod_bprop + # NOTE: environment variables in this function are read with + # os.environ.get, NOT os.getenv, on purpose: dynamo installs guards on + # os.environ reads (so changing an NVTE_* variable triggers recompilation + # under torch.compile), while os.getenv reads are unguarded and would bake + # stale values into compiled graphs. New code must follow suit. + # Run config - logger = logging.getLogger("DotProductAttention") - logger.setLevel(AttentionLogging._log_level) - if not logger.hasHandlers(): - logger.addHandler(AttentionLogging._stream_handler) + if torch.compiler.is_compiling(): + # logging.Logger methods graph-break under torch.compile; backend + # selection logs are only emitted in eager mode. + logger = _no_op_logger + else: + logger = logging.getLogger("DotProductAttention") + logger.setLevel(AttentionLogging._log_level) + if not logger.hasHandlers(): + logger.addHandler(AttentionLogging._stream_handler) device_compute_capability = get_device_compute_capability() cudnn_version = get_cudnn_version() run_config = { "transformer_engine_version": te.__version__, - "compute_capability": "sm" - + str(10 * device_compute_capability[0] + device_compute_capability[1]), + "compute_capability": ( + "sm" + str(10 * device_compute_capability[0] + device_compute_capability[1]) + ), "cuda_version": torch.version.cuda, "flash_attn_version": ( str(FlashAttentionUtils.version) @@ -429,27 +501,27 @@ def get_attention_backend( # Add FP8 environment variables to config if fp8: # all FP8 recipes: 1: (FP8 fwd, FP8 bwd), 0: (FP8 fwd, F16 bwd) - run_config["NVTE_FP8_DPA_BWD"] = int(os.getenv("NVTE_FP8_DPA_BWD", "1")) + run_config["NVTE_FP8_DPA_BWD"] = int(os.environ.get("NVTE_FP8_DPA_BWD", "1")) # Float8CurrentScaling: 1: use F16 O in bwd, 0: use FP8 O in bwd - run_config["NVTE_DPA_FP8CS_O_in_F16"] = int(os.getenv("NVTE_DPA_FP8CS_O_in_F16", "1")) + run_config["NVTE_DPA_FP8CS_O_in_F16"] = int(os.environ.get("NVTE_DPA_FP8CS_O_in_F16", "1")) # switch recipe to "F16", "DelayedScaling", or "Float8CurrentScaling" - _dpa_fp8_recipe = os.getenv("NVTE_DPA_FP8_RECIPE", "") + _dpa_fp8_recipe = os.environ.get("NVTE_DPA_FP8_RECIPE", "") run_config["NVTE_DPA_FP8_RECIPE"] = _dpa_fp8_recipe if _dpa_fp8_recipe != "": # config new recipe if switched - run_config["NVTE_DPA_FP8_FORMAT"] = os.getenv("NVTE_DPA_FP8_FORMAT", "HYBRID") - run_config["NVTE_DPA_FP8DS_AMAX_ALGO"] = os.getenv( + run_config["NVTE_DPA_FP8_FORMAT"] = os.environ.get("NVTE_DPA_FP8_FORMAT", "HYBRID") + run_config["NVTE_DPA_FP8DS_AMAX_ALGO"] = os.environ.get( "NVTE_DPA_FP8DS_AMAX_ALGO", "most_recent" ) run_config["NVTE_DPA_FP8DS_AMAX_HISTLEN"] = int( - os.getenv("NVTE_DPA_FP8DS_AMAX_HISTLEN", "1") + os.environ.get("NVTE_DPA_FP8DS_AMAX_HISTLEN", "1") ) run_config["NVTE_DPA_FP8DS_REDUCE_AMAX"] = int( - os.getenv("NVTE_DPA_FP8DS_REDUCE_AMAX", "1") + os.environ.get("NVTE_DPA_FP8DS_REDUCE_AMAX", "1") ) # UnfusedDotProductAttention: 1: allow FP8 emulation, 0: do not allow run_config["NVTE_UnfusedDPA_Emulate_FP8"] = int( - os.getenv("NVTE_UnfusedDPA_Emulate_FP8", "0") + os.environ.get("NVTE_UnfusedDPA_Emulate_FP8", "0") ) logger.debug("Running with config=%s", run_config) @@ -460,13 +532,13 @@ def get_attention_backend( qkv_format, q_format, kv_format = get_qkv_format(qkv_layout, inference_params) # Filter: Environment variables - use_flash_attention = int(os.getenv("NVTE_FLASH_ATTN", "1")) - use_flash_attention_2 = use_flash_attention and int(os.getenv("NVTE_FLASH_ATTN_V2", "1")) - use_flash_attention_3 = use_flash_attention and int(os.getenv("NVTE_FLASH_ATTN_V3", "1")) - use_flash_attention_4 = use_flash_attention and int(os.getenv("NVTE_FLASH_ATTN_V4", "1")) + use_flash_attention = int(os.environ.get("NVTE_FLASH_ATTN", "1")) + use_flash_attention_2 = use_flash_attention and int(os.environ.get("NVTE_FLASH_ATTN_V2", "1")) + use_flash_attention_3 = use_flash_attention and int(os.environ.get("NVTE_FLASH_ATTN_V3", "1")) + use_flash_attention_4 = use_flash_attention and int(os.environ.get("NVTE_FLASH_ATTN_V4", "1")) flash_attention_backend = None - use_fused_attention = int(os.getenv("NVTE_FUSED_ATTN", "1")) - use_unfused_attention = int(os.getenv("NVTE_UNFUSED_ATTN", "1")) + use_fused_attention = int(os.environ.get("NVTE_FUSED_ATTN", "1")) + use_unfused_attention = int(os.environ.get("NVTE_UNFUSED_ATTN", "1")) if not use_flash_attention_2 and FlashAttentionUtils.is_installed: logger.debug("Disabling FlashAttention 2 due to NVTE_FLASH_ATTN=0 or NVTE_FLASH_ATTN_V2=0") if not use_flash_attention_3 and FlashAttentionUtils.v3_is_installed: @@ -502,10 +574,10 @@ def _disable_all_flash_attention() -> None: if use_flash_attention_3 and FlashAttentionUtils.v3_is_installed: logger.debug("Disabling FlashAttention 3 for compute capability != sm90") use_flash_attention_3 = False - # FA4 supports SM80, SM90, SM100, SM120 - if not IS_HIP_EXTENSION and device_compute_capability < (8, 0): + # FA4 does not currently support SM8x. + if not IS_HIP_EXTENSION and device_compute_capability < (9, 0): if use_flash_attention_4 and FlashAttentionUtils.v4_is_installed: - logger.debug("Disabling FlashAttention 4 for compute capability < sm80") + logger.debug("Disabling FlashAttention 4 for compute capability < sm90") use_flash_attention_4 = False # On SM90, prefer FA3 over FA4 when FA3 is available. # FA3 is more mature on Hopper; FA4's SM90 backward has limitations @@ -591,7 +663,8 @@ def _disable_all_flash_attention() -> None: use_flash_attention_3 = False if use_unfused_attention: allow_emulation = ( - os.getenv("NVTE_UnfusedDPA_Emulate_FP8", "0") == "1" or is_in_onnx_export_mode() + os.environ.get("NVTE_UnfusedDPA_Emulate_FP8", "0") == "1" + or is_in_onnx_export_mode() ) if not allow_emulation: logger.debug("Disabling UnfusedDotProductAttention for FP8 attention") @@ -940,6 +1013,19 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt device_compute_capability[0] * 10 + device_compute_capability[1], ) use_flash_attention_4 = False + # FA4's validator currently accepts symmetric (512, 512) on SM100/SM110, + # but the generic forward kernel exceeds its TMEM allocation for that shape. + # Preserve the supported asymmetric (64, 512) MLA path while D512 support + # is completed upstream. + if ( + use_flash_attention_4 + and (10, 0) <= device_compute_capability < (12, 0) + and head_dim_qk == head_dim_v == 512 + ): + logger.debug( + "Disabling FlashAttention 4 for unsupported symmetric head_dim=512 on SM100/SM110." + ) + use_flash_attention_4 = False # flash-attn-4 4.0.0b11 validates (256, 256) on SM100, but its dedicated # hd256 kernel diverges from the reference for cross-attention/decode-like # shapes such as sq=1, skv=2048. Keep FA4 enabled for the self-attention @@ -1159,12 +1245,6 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt " bias for THD format" ) use_fused_attention = False - elif fp8 and fp8_meta["recipe"].fp8_dpa and qkv_format == "thd": - logger.debug( - "Disabling FusedAttention as it does not support context parallelism with FP8" - " attention and THD format" - ) - use_fused_attention = False elif fp8 and fp8_meta["recipe"].fp8_dpa and core_attention_bias_type != "no_bias": logger.debug( "Disabling FusedAttention as it does not support context parallelism with FP8" @@ -1178,13 +1258,6 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt cp_comm_type, ) use_fused_attention = False - elif qkv_format == "thd" and cp_comm_type in ["a2a+p2p"]: - logger.debug( - "Disabling FusedAttention as it does not support context parallelism with THD" - " format and cp_comm_type = %s", - cp_comm_type, - ) - use_fused_attention = False elif ( window_size is not None and (window_size[0] != -1 or window_size[1] not in [-1, 0]) @@ -1256,6 +1329,18 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt # | | converts window_size to an 'arbitrary' mask if window_size is None: window_size = check_set_window_size(attn_mask_type, window_size) + if ( + use_flash_attention_4 + and (10, 0) <= device_compute_capability < (12, 0) + and head_dim_qk == head_dim_v == 256 + and (window_size[0] != -1 or window_size[1] not in [-1, 0]) + ): + logger.debug( + "Disabling FlashAttention 4 as SM100 head_dim=256 does not support " + "sliding-window/local attention yet. Found: window_size = %s.", + window_size, + ) + use_flash_attention_4 = False if use_fused_attention and (window_size[0] != -1 or window_size[1] not in [-1, 0]): if ( fp8 @@ -1383,14 +1468,17 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt if fp8 and fp8_meta["recipe"].fp8_dpa: q_type = get_fp8_te_dtype(fp8_meta["recipe"], fprop_tensor=True) kv_type = q_type - fused_attention_backend = tex.get_fused_attn_backend( + # NOTE: under torch.compile the numeric args below must not be symbolic + # (assume_constant_result requires concrete values); ints/floats made + # dynamic by automatic dynamic currently graph break here. + fused_attention_backend = _get_fused_attn_backend( is_training, q_type, kv_type, - QKVLayout[qkv_layout], - AttnBiasType[fu_core_attention_bias_type], - AttnMaskType[attn_mask_type], - SoftmaxType[softmax_type], + qkv_layout, + fu_core_attention_bias_type, + attn_mask_type, + softmax_type, attention_dropout, num_heads, num_gqa_groups, @@ -1624,6 +1712,15 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt ) +@torch.no_grad() +def get_thd_padding_mask(num_tokens, cu_seqlens, cu_seqlens_padded): + """Identify inter-sequence padding in a flattened packed THD buffer.""" + rows = torch.arange(num_tokens, device=cu_seqlens_padded.device) + sequence = torch.searchsorted(cu_seqlens_padded[1:], rows, right=True) + valid_end = cu_seqlens_padded[sequence] + cu_seqlens[sequence + 1] - cu_seqlens[sequence] + return rows >= valid_end + + @torch.no_grad() def get_padding_mask( batch_size: int, @@ -1633,51 +1730,37 @@ def get_padding_mask( max_seqlen_kv: int = None, attention_type: str = "self", ): - """Convert cu_seqlens to attention_mask""" + """Convert cu_seqlens to attention_mask. + + Built with device-side ops only: reading the sequence lengths on the host + would synchronize the device once per sequence. + """ assert ( cu_seqlens_q is not None and max_seqlen_q is not None ), "cu_seqlens_q and max_seqlen_q are required for self-attention and cross-attention" - seqlens_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1] - attention_mask_q = torch.Tensor([]).to(dtype=torch.bool) - if attention_type == "cross": - assert ( - cu_seqlens_kv is not None and max_seqlen_kv is not None - ), "cu_seqlens_kv and max_seqlen_kv are required for cross-attention" - seqlens_kv = cu_seqlens_kv[1:] - cu_seqlens_kv[:-1] - attention_mask_kv = torch.Tensor([]).to(dtype=torch.bool) - for i in range(batch_size): - attention_mask_q = torch.cat( - [ - attention_mask_q, - torch.Tensor([False] * seqlens_q[i] + [True] * (max_seqlen_q - seqlens_q[i])) - .to(dtype=torch.bool) - .unsqueeze(0) - .unsqueeze(0) - .unsqueeze(0), - ], - dim=0, + + def _mask(cu_seqlens: torch.Tensor, max_seqlen: int) -> torch.Tensor: + # cu_seqlens may be longer than batch_size + 1 -- inference allocates it + # for the maximum batch size -- so only its first batch_size + 1 entries + # describe the current batch. + seqlens = cu_seqlens[1 : batch_size + 1] - cu_seqlens[:batch_size] + positions = torch.arange(max_seqlen, device=cu_seqlens.device) + # True marks a padding token, i.e. one beyond the sequence length. The + # mask is applied to the attention scores, so it goes on the device + # those live on -- a no-op unless cu_seqlens is a CPU tensor. + return ( + (positions.unsqueeze(0) >= seqlens.unsqueeze(1)) + .view(batch_size, 1, 1, max_seqlen) + .to(device="cuda") ) - if attention_type == "cross": - attention_mask_kv = torch.cat( - [ - attention_mask_kv, - torch.Tensor([False] * seqlens_kv[i] + [True] * (max_seqlen_kv - seqlens_kv[i])) - .to(dtype=torch.bool) - .unsqueeze(0) - .unsqueeze(0) - .unsqueeze(0), - ], - dim=0, - ) - attention_mask_q = attention_mask_q.to(device="cuda") + + attention_mask_q = _mask(cu_seqlens_q, max_seqlen_q) if attention_type == "self": - attention_mask = attention_mask_q - else: - attention_mask = ( - attention_mask_q, - attention_mask_kv.to(device="cuda"), - ) - return attention_mask + return attention_mask_q + assert ( + cu_seqlens_kv is not None and max_seqlen_kv is not None + ), "cu_seqlens_kv and max_seqlen_kv are required for cross-attention" + return attention_mask_q, _mask(cu_seqlens_kv, max_seqlen_kv) @torch.no_grad() @@ -2194,76 +2277,136 @@ def backward(ctx, grad_output): return None, None, _pack_tensor(indices, grad_output) -class ConvertTHDtoBSHD(torch.autograd.Function): +# --------------------------------------------------------------------------- +# THD <-> BSHD conversions exposed as `torch.library` custom ops so that +# `torch.compile` can trace them. Backward of each direction is the other +# direction, wired through `register_autograd` + `setup_context`, mirroring +# the pattern in `transformer_engine/pytorch/permutation.py`. +# --------------------------------------------------------------------------- + + +@torch.library.custom_op("te_attention::convert_thd_to_bshd", mutates_args=()) +def _convert_thd_to_bshd_op( + thd_tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + batch_size: int, + max_seqlen: int, +) -> torch.Tensor: + """Forward pass for THD->BSHD conversion.""" + if not thd_tensor.is_contiguous(): + thd_tensor = thd_tensor.contiguous() + return tex.convert_thd_to_bshd(thd_tensor, cu_seqlens, batch_size, max_seqlen) + + +@_convert_thd_to_bshd_op.register_fake +def _convert_thd_to_bshd_fake( + thd_tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + batch_size: int, + max_seqlen: int, +) -> torch.Tensor: + del cu_seqlens + h, d = thd_tensor.shape[1], thd_tensor.shape[2] + return torch.empty( + (batch_size, max_seqlen, h, d), dtype=thd_tensor.dtype, device=thd_tensor.device + ) + + +@torch.library.custom_op("te_attention::convert_bshd_to_thd", mutates_args=()) +def _convert_bshd_to_thd_op( + bshd_tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + num_tokens: int, +) -> torch.Tensor: + """Forward pass for BSHD->THD conversion.""" + if not bshd_tensor.is_contiguous(): + bshd_tensor = bshd_tensor.contiguous() + return tex.convert_bshd_to_thd(bshd_tensor, cu_seqlens, num_tokens) + + +@_convert_bshd_to_thd_op.register_fake +def _convert_bshd_to_thd_fake( + bshd_tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + num_tokens: int, +) -> torch.Tensor: + del cu_seqlens + h, d = bshd_tensor.shape[2], bshd_tensor.shape[3] + return torch.empty((num_tokens, h, d), dtype=bshd_tensor.dtype, device=bshd_tensor.device) + + +def _convert_thd_to_bshd_setup_context(ctx, inputs, output): + del output + thd_tensor, cu_seqlens, _batch_size, _max_seqlen = inputs + ctx.save_for_backward(cu_seqlens) + ctx.num_tokens = thd_tensor.size(0) + + +def _convert_thd_to_bshd_backward_wrapper(ctx, grad_bshd): + (cu_seqlens,) = ctx.saved_tensors + grad_thd = torch.ops.te_attention.convert_bshd_to_thd(grad_bshd, cu_seqlens, ctx.num_tokens) + return grad_thd, None, None, None + + +_convert_thd_to_bshd_op.register_autograd( + _convert_thd_to_bshd_backward_wrapper, + setup_context=_convert_thd_to_bshd_setup_context, +) + + +def _convert_bshd_to_thd_setup_context(ctx, inputs, output): + del output + bshd_tensor, cu_seqlens, _num_tokens = inputs + ctx.save_for_backward(cu_seqlens) + ctx.batch_size = bshd_tensor.size(0) + ctx.max_seqlen = bshd_tensor.size(1) + + +def _convert_bshd_to_thd_backward_wrapper(ctx, grad_thd): + (cu_seqlens,) = ctx.saved_tensors + grad_bshd = torch.ops.te_attention.convert_thd_to_bshd( + grad_thd, cu_seqlens, ctx.batch_size, ctx.max_seqlen + ) + return grad_bshd, None, None + + +_convert_bshd_to_thd_op.register_autograd( + _convert_bshd_to_thd_backward_wrapper, + setup_context=_convert_bshd_to_thd_setup_context, +) + + +class ConvertTHDtoBSHD: """ Convert a tensor from qkv_format = thd to qkv_format = bshd. + + Thin wrapper around the ``te_attention::convert_thd_to_bshd`` custom op, + exposing an ``.apply(thd_tensor, cu_seqlens, max_seqlen)`` staticmethod so + callsites keep the ``autograd.Function``-style ``.apply(...)`` invocation. """ @staticmethod - def forward(ctx, thd_tensor, cu_seqlens, max_seqlen): + def apply(thd_tensor, cu_seqlens, max_seqlen): # pylint: disable=missing-function-docstring batch_size = cu_seqlens.shape[0] - 1 - if not thd_tensor.is_contiguous(): - thd_tensor = thd_tensor.contiguous() - bshd_tensor = tex.convert_thd_to_bshd( - thd_tensor, - cu_seqlens, - batch_size, - max_seqlen, - ) - ctx.save_for_backward(cu_seqlens) - ctx.num_tokens = thd_tensor.shape[0] - return bshd_tensor - - @staticmethod - def backward(ctx, bshd_tensor): - # pylint: disable=missing-function-docstring - (cu_seqlens,) = ctx.saved_tensors - if not bshd_tensor.is_contiguous(): - bshd_tensor = bshd_tensor.contiguous() - thd_tensor = tex.convert_bshd_to_thd( - bshd_tensor, - cu_seqlens, - ctx.num_tokens, + return torch.ops.te_attention.convert_thd_to_bshd( + thd_tensor, cu_seqlens, batch_size, max_seqlen ) - return thd_tensor, None, None -class ConvertBSHDtoTHD(torch.autograd.Function): +class ConvertBSHDtoTHD: """ Convert a tensor from qkv_format = bshd to qkv_format = thd. - """ - @staticmethod - def forward(ctx, bshd_tensor, cu_seqlens): - # pylint: disable=missing-function-docstring - num_tokens = cu_seqlens[-1] - max_seqlen = bshd_tensor.shape[1] - if not bshd_tensor.is_contiguous(): - bshd_tensor = bshd_tensor.contiguous() - thd_tensor = tex.convert_bshd_to_thd( - bshd_tensor, - cu_seqlens, - num_tokens, - ) - ctx.save_for_backward(cu_seqlens) - ctx.max_seqlen = max_seqlen - return thd_tensor + Thin wrapper around the ``te_attention::convert_bshd_to_thd`` custom op, + exposing an ``.apply(bshd_tensor, cu_seqlens, num_tokens)`` staticmethod so + callsites keep the ``autograd.Function``-style ``.apply(...)`` invocation. + """ @staticmethod - def backward(ctx, thd_tensor): + def apply(bshd_tensor, cu_seqlens, num_tokens): # pylint: disable=missing-function-docstring - (cu_seqlens,) = ctx.saved_tensors - batch_size = cu_seqlens.shape[0] - 1 - if not thd_tensor.is_contiguous(): - thd_tensor = thd_tensor.contiguous() - bshd_tensor = tex.convert_thd_to_bshd( - thd_tensor, - cu_seqlens, - batch_size, - ctx.max_seqlen, - ) - return bshd_tensor, None + return torch.ops.te_attention.convert_bshd_to_thd(bshd_tensor, cu_seqlens, num_tokens) def get_qkv_format( @@ -2484,6 +2627,22 @@ def run_iteratively(q, k, v): if qkv_layout == "not_supported": raise RuntimeError("The provided qkv memory layout is not supported!") + if len(qkv_layout.split("_")) < 3: + # q/k/v were recognized as views of a packed buffer only by inspecting + # their data pointers, strides and storage offsets. Skip the nudge while + # CPU offloading is enabled: offloading forces MultiheadAttention onto + # its sliced-views fallback, so packed views reaching detection are + # expected there and the caller has no migration option. + if not is_cpu_offload_enabled(): + warnings.warn( + "Relying on pointer-based detection of packed q/k/v layouts" + f" (detected {qkv_layout!r}) is deprecated: pass the packed buffer" + " explicitly via qkv_layer/kv_layer (with qkv_interleave_dim) to" + " DotProductAttention instead.", + DeprecationWarning, + stacklevel=2, + ) + if inference_params is not None and inference_params.is_paged: qkv_layout = "paged_kv_" + qkv_layout @@ -2593,14 +2752,15 @@ def get_attention_quantizers(fp8, quantizers): ]: if _q is None and _name in _allow_none: continue - assert isinstance(_q, _fp8_types), ( - "FP8 attention requires FP8-compatible quantizers for all DPA tensor slots, " - f"but {_name} quantizer is {type(_q).__name__}. " - "When using CustomRecipe with fp8_dpa=True, ensure the factory returns an " - "FP8 quantizer (Float8Quantizer, Float8CurrentScalingQuantizer, or " - "MXFP8Quantizer) for all DPA roles (module_type='dpa') and for None roles " - "(boundary slots like O output and dQKV grad-input)." - ) + if not isinstance(_q, _fp8_types): + raise TypeError( + "FP8 attention requires FP8-compatible quantizers for all DPA tensor slots, " + f"but {_name} quantizer is {type(_q).__name__}. " + "When using CustomRecipe with fp8_dpa=True, ensure the factory returns an " + "FP8 quantizer (Float8Quantizer, Float8CurrentScalingQuantizer, or " + "MXFP8Quantizer) for all DPA roles (module_type='dpa') and for None roles " + "(boundary slots like O output and dQKV grad-input)." + ) return QKV_quantizer, O_quantizer, S_quantizer, dQKV_quantizer, dO_quantizer, dP_quantizer @@ -2699,10 +2859,37 @@ def mxfp8_quantize_fast_path(tensor_quantizer_pairs, src_format): """ if not tensor_quantizer_pairs: return [], src_format + + fp8_tensors = mxfp8_quantize_only(tensor_quantizer_pairs, src_format) + mxfp8_transpose_swizzle(fp8_tensors, src_format) + return fp8_tensors, "bhsd" + + +def mxfp8_quantize_only(tensor_quantizer_pairs, src_format): + """Phase 1 of mxfp8_quantize_fast_path: quantize only, no BHSD transpose or GEMM swizzle. + + Returns MXFP8Tensors with data and scale_invs reshaped to src_format layout. + Call mxfp8_transpose_swizzle to complete the BHSD permute + swizzle when ready + (e.g. after pre-quantized tensors from fused kernels are also available). + + Parameters + ---------- + tensor_quantizer_pairs : list of (torch.Tensor, MXFP8Quantizer) + Same contract as mxfp8_quantize_fast_path. + src_format : str + ``"bshd"`` or ``"sbhd"``. + + Returns + ------- + fp8_tensors : list of MXFP8Tensor + Data and scale_invs in src_format layout; NOT yet BHSD-permuted or swizzled. + """ + if not tensor_quantizer_pairs: + return [] assert src_format in ( "bshd", "sbhd", - ), f"mxfp8_quantize_fast_path only supports bshd/sbhd, got {src_format!r}." + ), f"mxfp8_quantize_only only supports bshd/sbhd, got {src_format!r}." _s_dim = {"bshd": 1, "sbhd": 0} _d_dim = {"bshd": 3, "sbhd": 3} @@ -2713,45 +2900,74 @@ def mxfp8_quantize_fast_path(tensor_quantizer_pairs, src_format): rs_shape[_d_dim[src_format]] //= MXFP8_BLOCK_SCALING_SIZE cs_shape = list(original_shape) cs_shape[_s_dim[src_format]] //= MXFP8_BLOCK_SCALING_SIZE - - # view tensor as 2D for quantization - # BSHD -> (B*S, H*D) - # SBHD -> (S, B*H*D) if src_format == "bshd": - tensor = tensor.view(*tensor.shape[:2], -1) + t2d = tensor.view(*tensor.shape[:2], -1) else: - tensor = tensor.view(tensor.shape[0], -1) - - # quantize + t2d = tensor.view(tensor.shape[0], -1) orig_optimize = quantizer.optimize_for_gemm quantizer.optimize_for_gemm = False - fp8_tensor = quantizer(tensor) + fp8_2d = quantizer(t2d) quantizer.optimize_for_gemm = orig_optimize - - # reshape rowwise/columnwise data to original shape - fp8_tensor._rowwise_data = ( - fp8_tensor._rowwise_data.view(original_shape) - if fp8_tensor._rowwise_data is not None - else None - ) - fp8_tensor._columnwise_data = ( - fp8_tensor._columnwise_data.view(original_shape) - if fp8_tensor._columnwise_data is not None - else None - ) - fp8_tensor._rowwise_scale_inv = ( - fp8_tensor._rowwise_scale_inv.view(rs_shape) - if fp8_tensor._rowwise_scale_inv is not None - else None - ) - fp8_tensor._columnwise_scale_inv = ( - fp8_tensor._columnwise_scale_inv.view(cs_shape) - if fp8_tensor._columnwise_scale_inv is not None - else None + # Re-wrap with the original 4D SBHD/BSHD shape so that shape[-1] equals the per-head + # dimension (matching Q's wrapper shape) and fused_attn_bwd produces 4D dkv that + # matches key/value's expected gradient shape in _KFQuantizeKVForAttn.backward. + fp8_t = MXFP8Tensor( + shape=original_shape, + dtype=tensor.dtype, + rowwise_data=( + fp8_2d._rowwise_data.view(original_shape) + if fp8_2d._rowwise_data is not None + else None + ), + rowwise_scale_inv=( + fp8_2d._rowwise_scale_inv.view(rs_shape) + if fp8_2d._rowwise_scale_inv is not None + else None + ), + columnwise_data=( + fp8_2d._columnwise_data.view(original_shape) + if fp8_2d._columnwise_data is not None + else None + ), + columnwise_scale_inv=( + fp8_2d._columnwise_scale_inv.view(cs_shape) + if fp8_2d._columnwise_scale_inv is not None + else None + ), + quantizer=quantizer, + requires_grad=False, + fp8_dtype=fp8_2d._fp8_dtype, + with_gemm_swizzled_scales=False, ) - fp8_tensors.append(fp8_tensor) + fp8_tensors.append(fp8_t) + return fp8_tensors + + +def mxfp8_transpose_swizzle(fp8_tensors, src_format): + """Phase 2 of mxfp8_quantize_fast_path: batched BHSD-transpose + GEMM-swizzle. + + For tensors whose data is already quantized (e.g. from a fused GEMM+quant kernel + or from mxfp8_quantize_only), permutes each tensor's scale_invs from src_format to + BHSD and applies the GEMM swizzle in-place. Complements mxfp8_quantize_only to + allow pre-quantized tensors (like a fused-kernel Q) to be processed in the same + batched operation as freshly quantized K/V. + + Parameters + ---------- + fp8_tensors : list of MXFP8Tensor + Tensors with _rowwise_scale_inv / _columnwise_scale_inv in src_format layout. + Modified in-place: scale_invs are replaced with BHSD-permuted, swizzled versions. + src_format : str + ``"bshd"`` or ``"sbhd"``. + """ + if not fp8_tensors: + return + + assert src_format in ( + "bshd", + "sbhd", + ), f"mxfp8_transpose_swizzle only supports bshd/sbhd, got {src_format!r}." - # ---- Pad + permute + swizzle scale_inv to BHSD ---- rs_list = [t._rowwise_scale_inv for t in fp8_tensors] cs_list = [t._columnwise_scale_inv for t in fp8_tensors] @@ -2785,50 +3001,25 @@ def _build_outputs(scale_list, alignment): buf = torch.empty(total, dtype=torch.uint8, device=device) return [buf[e[0] : e[0] + e[1]].view(e[2]) if e is not None else None for e in entries] - # allocate buffers with padding in mind rs_outs = _build_outputs(rs_list, 4) cs_outs = _build_outputs(cs_list, 128) - # permute scale_invs to BHSD; batched rs_permuted = tex.multi_tensor_transpose_to_bhsd( - rs_list, - original_format=src_format, - outputs=rs_outs, + rs_list, original_format=src_format, outputs=rs_outs ) cs_permuted = tex.multi_tensor_transpose_to_bhsd( - cs_list, - original_format=src_format, - outputs=cs_outs, + cs_list, original_format=src_format, outputs=cs_outs ) - # build output tensors - result = [] for t, rp, cp in zip(fp8_tensors, rs_permuted, cs_permuted): - rp = rp.view(-1, rp.shape[-1]) if rp is not None else None - cp = cp.view(-1, cp.shape[-1]) if cp is not None else None - result.append( - MXFP8Tensor( - shape=t.shape, - dtype=t.dtype, - rowwise_data=t._rowwise_data, - rowwise_scale_inv=rp, - columnwise_data=t._columnwise_data, - columnwise_scale_inv=cp, - quantizer=t._quantizer, - requires_grad=False, - fp8_dtype=t._fp8_dtype, - with_gemm_swizzled_scales=t._with_gemm_swizzled_scales, - ) - ) + t._rowwise_scale_inv = rp.view(-1, rp.shape[-1]) if rp is not None else None + t._columnwise_scale_inv = cp.view(-1, cp.shape[-1]) if cp is not None else None - # swizzle in place; batched - tex.multi_tensor_swizzle_scales_for_gemm_unchecked_(result, True, False) - tex.multi_tensor_swizzle_scales_for_gemm_unchecked_(result, False, True) - for t in result: + tex.multi_tensor_swizzle_scales_for_gemm_unchecked_(fp8_tensors, True, False) + tex.multi_tensor_swizzle_scales_for_gemm_unchecked_(fp8_tensors, False, True) + for t in fp8_tensors: t._with_gemm_swizzled_scales = True - return result, "bhsd" - def combine_and_quantize( qkv_layout, @@ -2839,8 +3030,18 @@ def combine_and_quantize( used_in_forward=True, used_in_backward=False, keep_same_data_and_scale_inv_format=False, + combined_qkv: Optional[torch.Tensor] = None, + combined_kv: Optional[torch.Tensor] = None, ): - """Combine Q, K, V tensors based on qkv_layout and quantize them together.""" + """Combine Q, K, V tensors based on qkv_layout and quantize them together. + + When ``combined_qkv`` (for ``qkv_group=1`` layouts such as ``bs3hd``) or + ``combined_kv`` (for ``qkv_group=2`` layouts such as ``bshd_bs2hd``) is provided, it must be the + caller's original packed buffer that q/k/v are views of. It is then quantized + directly instead of re-deriving the packed buffer from the q/k/v views via + ``combine_tensors`` (which rebuilds it with a raw ``set_`` under a silent + adjacency/interleave assumption). Ignored for MXFP8 quantization. + """ if isinstance(qkv_quantizer, MXFP8Quantizer): qkv_format, q_format, kv_format = get_qkv_format(qkv_layout) assert qkv_format in ("bshd", "sbhd"), ( @@ -2940,12 +3141,26 @@ def combine_and_quantize( match qkv_group: case 1: dim = qkv_layout.find("3") - qkv = combine_tensors([q, k, v], dim) + if combined_qkv is not None: + assert combined_qkv.shape[dim] == 3, ( + f"combined_qkv does not match qkv_layout {qkv_layout}: expected" + f" size 3 at dim {dim}, got shape {tuple(combined_qkv.shape)}." + ) + qkv = combined_qkv + else: + qkv = combine_tensors([q, k, v], dim) qkv_fp8 = qkv_quantizer(qkv) q_data, k_data, v_data = SplitAlongDim.apply(qkv_fp8._data, dim, [1, 1, 1], True) case 2: dim = qkv_layout.split("_")[1].find("2") - kv = combine_tensors([k, v], dim) + if combined_kv is not None: + assert combined_kv.shape[dim] == 2, ( + f"combined_kv does not match qkv_layout {qkv_layout}: expected" + f" size 2 at dim {dim}, got shape {tuple(combined_kv.shape)}." + ) + kv = combined_kv + else: + kv = combine_tensors([k, v], dim) tensors = [q, kv] num_tensors = len(tensors) shapes = [x.shape for x in tensors] diff --git a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py new file mode 100644 index 0000000000..c176985254 --- /dev/null +++ b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py @@ -0,0 +1,175 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fused MLA Q up-projection + per-head RoPE + MXFP8 quantize.""" + +from __future__ import annotations +import functools +import os +from importlib.metadata import PackageNotFoundError, version as get_pkg_version + +import torch +import transformer_engine_torch as tex +from packaging.version import Version as PkgVersion + +from ..constants import MXFP8_BLOCK_SCALING_SIZE +from ..quantized_tensor import QuantizedTensor +from ..tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor +from ..utils import get_device_compute_capability + +_CUDNN_FRONTEND_MIN_VERSION = "1.27.0" + + +def _cudnn_frontend_version_supported() -> bool: + """Check that the installed nvidia-cudnn-frontend meets the minimum version.""" + try: + return PkgVersion(get_pkg_version("nvidia-cudnn-frontend")) >= PkgVersion( + _CUDNN_FRONTEND_MIN_VERSION + ) + except PackageNotFoundError: + return False + + +class FusedMLAQUpProjRopeQuant: + """Wrapper for the cuDNN fused MLA Q up-proj + per-head RoPE + MXFP8 quantize kernel. + + - If w is already a QuantizedTensor (primary FP8 parameter in MXFP8BlockScaling recipe), + this performs an MXFP8 GEMM within the fusion (and quantizes the input if necessary) + - Otherwise (plain BF16 weight), x and w are passed as-is to the BF16 kernel variant. + """ + + @classmethod + @functools.lru_cache(maxsize=None) + def _kernel(cls): + # Import directly from the subpackage to avoid depending on cudnn/__init__.py + # lazy-import registration (which would require overlaying cudnn/__init__.py and + # could revert atomicrmw fixes present in the container's version). + try: + from cudnn import gemm_proj_rope_mxfp8_wrapper_sm100 + + return gemm_proj_rope_mxfp8_wrapper_sm100 + except ImportError: + return None + + @classmethod + @functools.lru_cache(maxsize=None) + def is_supported(cls) -> bool: + """Whether the cuDNN FE fused gemm rope quant wrapper is available""" + if int(os.environ.get("NVTE_FUSED_MLA_Q_UPROJ", "1")) <= 0: + return False + if not _cudnn_frontend_version_supported(): + return False + if get_device_compute_capability()[0] < 10: + return False + if cls._kernel() is None: + return False + return True + + @classmethod + def run( + cls, + x: torch.Tensor, + w, # MXFP8Tensor (primary FP8 param) or bf16 torch.Tensor + cos: torch.Tensor, + sin: torch.Tensor, + s: int, + b: int, + ) -> "tuple[MXFP8Tensor, torch.Tensor]": + """Run the fused kernel; return (Q MXFP8Tensor, activation saved for the wgrad backward). + + The kernel precision is selected by the weight precision. + """ + + from cuda.bindings import driver as cuda + + stream = cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + wrapper = cls._kernel() + + if isinstance(w, QuantizedTensor): + assert isinstance(w, MXFP8Tensor), ( + "FusedMLAQUpProjRopeQuant expects an MXFP8Tensor weight (MXFP8BlockScaling" + f" recipe), got {type(w).__name__}. Use the unfused path for other quantization" + " recipes." + ) + # ---- FP8 projection: MXFP8-cast x (both usages) + reuse w's fp8 codes -> mxfp8in ---- + # Quantize x with both rowwise (for the forward GEMM) and columnwise (for the FP8 + # wgrad in backward, matching the unfused path). + x_quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True + ) + x_mxfp8 = x_quantizer(x) + x_code = x_mxfp8._rowwise_data.view(torch.float8_e4m3fn) # [tokens, K] + x_scale = x_mxfp8._rowwise_scale_inv # [tokens, K//32] uint8 + + # Primary FP8 parameter: already quantized; use its rowwise FP8 codes + E8M0 scales. + w.update_usage(rowwise_usage=True, columnwise_usage=None) + w_code = w._rowwise_data.view(torch.float8_e4m3fn) # [N, K] + w_scale = w._rowwise_scale_inv # [N, K//32] uint8 + + out = wrapper( + x_code, + w_code, + cos, + sin, + x_scale=x_scale, + w_scale=w_scale, + w_out_in=True, + stream=stream, + ) + + # Drop rowwise data now. + # Only columnwise x is needed for the FP8 wgrad in backward. + x_mxfp8.update_usage(rowwise_usage=False, columnwise_usage=True) + x_saved = x_mxfp8 + else: + # ---- 16-bit projection: bf16 GEMM inputs -> bf16in (the projection stays bf16) ---- + out = wrapper(x, w, cos, sin, w_out_in=True, stream=stream) + x_saved = x + + nh = out["out_fp8_row"].shape[1] + d = out["out_fp8_row"].shape[2] + query = cls.wrap_mxfp8( + out["out_fp8_row"], + out["out_scales_row"], + out["out_fp8_col"], + out["out_scales_col"], + s, + b, + nh, + d, + ) + # 2nd return is the activation to save for wgrad: MXFP8 (fp8 path) or bf16 (16-bit path). + return query, x_saved + + @classmethod + def wrap_mxfp8( + cls, + fp8_row: torch.Tensor, + scales_row: torch.Tensor, + fp8_col: torch.Tensor, + scales_col: torch.Tensor, + s: int, + b: int, + nh: int, + d: int, + ) -> MXFP8Tensor: + """Wrap raw data and scale tensors into an MXFP8Tensor""" + + blk = MXFP8_BLOCK_SCALING_SIZE + # Both rowwise and columnwise Q are required: + # - Forward QK^T uses rowwise + # - cuDNN backward (fused_attn_fp8_bwd_impl) requires columnwise for dK gradient + quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True) + return MXFP8Tensor( + shape=(s, b, nh, d), + dtype=torch.bfloat16, + rowwise_data=fp8_row.view(s, b, nh, d), + rowwise_scale_inv=scales_row.view(s, b, nh, d // blk), + columnwise_data=fp8_col.view(s, b, nh, d), + columnwise_scale_inv=scales_col.view(s // blk, b, nh, d), + quantizer=quantizer, + requires_grad=False, + fp8_dtype=tex.DType.kFloat8E4M3, + with_gemm_swizzled_scales=False, + ) diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index 70ae9dfc21..f87365bf7c 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -9,6 +9,7 @@ import torch from transformer_engine.pytorch.quantization import FP8GlobalStateManager, QuantizerRole +from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.module import LayerNormLinear, Linear, RMSNorm, LayerNorm @@ -499,11 +500,11 @@ def _update_output_quantizer_roles( 1. ``qkv_fp8_output`` — **QKV linear → DPA (fwd)**: the QKV linear's ``output_quantizer_role`` is told its consumer is DPA. - 2. ``proj_fp8_grad`` — **Proj linear ← DPA (bwd)**: proj's - ``grad_input_quantizer_role`` is told its producer is DPA. + 2. ``proj_fp8_grad`` — **Proj linear → DPA (bwd)**: proj's + ``grad_input_quantizer_role`` is told its consumer is DPA. 3. ``dpa_fp8_output`` — **DPA → Proj linear (fwd)**: DPA's ``output_quantizer_role`` is told its consumer is the proj linear. - 4. ``dpa_fp8_output`` — **DPA ← QKV linear (bwd)**: DPA's + 4. ``dpa_fp8_output`` — **DPA → QKV linear (bwd)**: DPA's ``grad_input_quantizer_role`` is told its consumer is QKV linear. When a flag is ``False`` the corresponding role is reset to ``None`` @@ -529,7 +530,7 @@ def _update_output_quantizer_roles( self.query_layer.output_quantizer_role = qkv_output_role self.key_value.output_quantizer_role = qkv_output_role - # ── Boundary 2 (bwd): Proj grad-input ← produced by DPA ────────── + # ── Boundary 2 (bwd): Proj grad-input (dO) → consumed by DPA ───── proj_grad_input_role = ( QuantizerRole(module_type="dpa", tensor_type="do", name=dpa_name) if proj_fp8_grad @@ -869,12 +870,27 @@ def forward( # ====================== fp8 = FP8GlobalStateManager.is_fp8_enabled() + custom_recipe = False if _dpa_fp8_recipe == "": fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() + custom_recipe = fp8_recipe.custom() fp8_dpa = fp8_recipe.fp8_dpa fp8_mha = fp8_recipe.fp8_mha float8_current_scaling = fp8_recipe.float8_current_scaling() mxfp8_scaling = fp8_recipe.mxfp8() + if fp8 and custom_recipe and fp8_mha: + # Wire every boundary this CustomRecipe may quantize before + # DPA materializes its recipe state. Some quantizer families + # disable the corresponding output below, but pre-wiring avoids + # rebuilding that state after inspecting its canonical QKV slot. + self._update_output_quantizer_roles( + rotary_pos_emb is None, + True, + True, + ) + float8_current_scaling, mxfp8_scaling = ( + self.core_attention.get_qkv_quantization_capabilities() + ) else: fp8_dpa = _dpa_fp8_recipe_dpa fp8_mha = _dpa_fp8_recipe_mha @@ -896,12 +912,35 @@ def forward( # DPA: produce FP8 output to take advantage of O amax from DPA; Projection Gemm can take FP8 or F16 inputs # 1. FP8DS/FP8CS recipe: produce FP8 output # 2. MXFP8 recipe: produce F16 output; again, due to quantization dimensions mismatch - dpa_fp8_output = fp8 and (fp8_dpa or fp8_mha) and not mxfp8_scaling + # For CustomRecipe, fp8_dpa only controls DPA-internal quantization. + # External MHA boundary tensors become FP8 only when fp8_mha is enabled. + dpa_fp8_output_enabled = fp8_mha if custom_recipe else (fp8_dpa or fp8_mha) + dpa_fp8_output = fp8 and dpa_fp8_output_enabled and not mxfp8_scaling # Projection Gemm: match DPA output except # 1. FP8CS recipe: produce F16 grads; again, due to cuBLAS limitation proj_fp8_grad = dpa_fp8_output and not float8_current_scaling - self._update_output_quantizer_roles(qkv_fp8_output, proj_fp8_grad, dpa_fp8_output) + # Custom fp8_mha boundaries were wired before DPA recipe-state setup so + # querying its canonical QKV quantizer cannot trigger a second build. + if not (fp8 and custom_recipe and fp8_mha and _dpa_fp8_recipe == ""): + self._update_output_quantizer_roles(qkv_fp8_output, proj_fp8_grad, dpa_fp8_output) + + # Packed pass-through to DotProductAttention: the fused QKV/KV projection + # already produces one packed buffer, which DPA accepts directly via its + # declarative qkv_layer/kv_layer arguments (deriving q/k/v as zero-copy + # views and skipping pointer-based layout detection). Only possible when + # no per-tensor operation (RoPE, QK normalization, KV caching, CPU + # offloading) needs the individual q/k/v slices. + packed_dpa_eligible = ( + rotary_pos_emb is None + and self.q_norm is None + and self.k_norm is None + and inference_params is None + and not is_cpu_offload_enabled() + ) + packed_qkv_layer = None + packed_kv_layer = None + packed_interleave_dim = -3 layernorm_output = None if self.attention_type == "self": @@ -947,28 +986,43 @@ def forward( mixed_x_layer = mixed_x_layer.view(*new_tensor_shape) - # qkv_weight_interleaved: - # [sq, b, ng, (np/ng + 2), hn] - # --> [sq, b, ng, np/ng, hn], [sq, b, ng, 1, hn], [sq, b, ng, 1, hn] - # not qkv_weight_interleaved: - # [sq, b, (np/ng + 2), ng, hn] - # --> [sq, b, np/ng, np, hn], [sq, b, 1, ng, hn], [sq, b, 1, ng, hn] - query_layer, key_layer, value_layer = SplitAlongDim.apply( - mixed_x_layer, split_dim, (num_queries_per_key_value, 1, 1) - ) - - if self.qkv_format == "thd": - query_layer, key_layer, value_layer = ( - x.reshape(x.size(0), -1, self.hidden_size_per_attention_head) - for x in (query_layer, key_layer, value_layer) - ) + if ( + num_queries_per_key_value == 1 + and packed_dpa_eligible + and not isinstance(mixed_x_layer, QuantizedTensorStorage) + ): + # np == ng: the projection output is a uniform 3-interleave + # ([.., h, 3, d] interleaved / [.., 3, h, d] otherwise), which + # DotProductAttention accepts directly as a declared packed + # qkv_layer -- no slicing here, no layout detection there. + packed_qkv_layer = mixed_x_layer + packed_interleave_dim = split_dim + query_layer = None + key_layer = None + value_layer = None else: - # query: -> [sq, b, np, hn] - # key, value: -> [sq, b, ng, hn] - query_layer, key_layer, value_layer = ( - x.reshape(x.size(0), x.size(1), -1, self.hidden_size_per_attention_head) - for x in (query_layer, key_layer, value_layer) + # qkv_weight_interleaved: + # [sq, b, ng, (np/ng + 2), hn] + # --> [sq, b, ng, np/ng, hn], [sq, b, ng, 1, hn], [sq, b, ng, 1, hn] + # not qkv_weight_interleaved: + # [sq, b, (np/ng + 2), ng, hn] + # --> [sq, b, np/ng, np, hn], [sq, b, 1, ng, hn], [sq, b, 1, ng, hn] + query_layer, key_layer, value_layer = SplitAlongDim.apply( + mixed_x_layer, split_dim, (num_queries_per_key_value, 1, 1) ) + + if self.qkv_format == "thd": + query_layer, key_layer, value_layer = ( + x.reshape(x.size(0), -1, self.hidden_size_per_attention_head) + for x in (query_layer, key_layer, value_layer) + ) + else: + # query: -> [sq, b, np, hn] + # key, value: -> [sq, b, ng, hn] + query_layer, key_layer, value_layer = ( + x.reshape(x.size(0), x.size(1), -1, self.hidden_size_per_attention_head) + for x in (query_layer, key_layer, value_layer) + ) elif self.attention_type == "cross": # Attention heads [sk, b, h] --> [sk, b, (ng * 2 * hn)] mixed_kv_layer = self.key_value( @@ -996,34 +1050,56 @@ def forward( mixed_kv_layer = mixed_kv_layer.view(*new_tensor_shape) - # mixed_kv_layer --> 2 [sk, b, ng, hn] - key_layer, value_layer = SplitAlongDim.apply( - mixed_kv_layer, - split_dim, - mixed_kv_layer.shape[split_dim] // 2, - ) - key_layer, value_layer = ( - x.reshape( - x.size(0), - x.size(1), - -1, - self.hidden_size_per_attention_head, - ) - for x in (key_layer, value_layer) - ) - - if self.qkv_format == "thd": - key_layer, value_layer = ( - x.reshape(x.size(0), -1, self.hidden_size_per_attention_head) - for x in (key_layer, value_layer) - ) + if packed_dpa_eligible and not isinstance(mixed_kv_layer, QuantizedTensorStorage): + # Declare the packed KV to DotProductAttention instead of + # slicing it: expose the 2-interleave as its own dimension. + if self.qkv_weight_interleaved: + # [.., ng, 2 * hn] --> [.., ng, 2, hn] + packed_kv_shape = mixed_kv_layer.size()[:-1] + ( + 2, + self.hidden_size_per_attention_head, + ) + packed_interleave_dim = -2 + else: + # [.., 2 * ng, hn] --> [.., 2, ng, hn] + packed_kv_shape = mixed_kv_layer.size()[:-2] + ( + 2, + self.num_gqa_groups_per_partition, + self.hidden_size_per_attention_head, + ) + packed_interleave_dim = -3 + packed_kv_layer = mixed_kv_layer.view(*packed_kv_shape) + key_layer = None + value_layer = None else: - # key, value: -> [sq, b, ng, hn] + # mixed_kv_layer --> 2 [sk, b, ng, hn] + key_layer, value_layer = SplitAlongDim.apply( + mixed_kv_layer, + split_dim, + mixed_kv_layer.shape[split_dim] // 2, + ) key_layer, value_layer = ( - x.reshape(x.size(0), x.size(1), -1, self.hidden_size_per_attention_head) + x.reshape( + x.size(0), + x.size(1), + -1, + self.hidden_size_per_attention_head, + ) for x in (key_layer, value_layer) ) + if self.qkv_format == "thd": + key_layer, value_layer = ( + x.reshape(x.size(0), -1, self.hidden_size_per_attention_head) + for x in (key_layer, value_layer) + ) + else: + # key, value: -> [sq, b, ng, hn] + key_layer, value_layer = ( + x.reshape(x.size(0), x.size(1), -1, self.hidden_size_per_attention_head) + for x in (key_layer, value_layer) + ) + # Attention head [sq, b, h] --> [sq, b, hp] if self.input_layernorm: layernorm_query_outputs = self.layernorm_query( @@ -1143,6 +1219,9 @@ def forward( inference_params=inference_params, pad_between_seqs=pad_between_seqs, fp8_output=dpa_fp8_output, + qkv_layer=packed_qkv_layer, + kv_layer=packed_kv_layer, + qkv_interleave_dim=packed_interleave_dim, ) # =================== diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 5a9226b580..74896e6e6d 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -7,6 +7,7 @@ """Python interface for fused attention extensions""" import math +from enum import IntEnum from typing import Tuple, List, Union, Optional import torch from torch.utils.cpp_extension import IS_HIP_EXTENSION @@ -100,19 +101,73 @@ "learnable": NVTE_Softmax_Type.NVTE_LEARNABLE_SOFTMAX, } -if not IS_HIP_EXTENSION: - FusedAttnBackend = { - "F16_max512_seqlen": NVTE_Fused_Attn_Backend.NVTE_F16_max512_seqlen, - "F16_arbitrary_seqlen": NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, - "FP8": NVTE_Fused_Attn_Backend.NVTE_FP8, - "No_Backend": NVTE_Fused_Attn_Backend.NVTE_No_Backend, - } -else: - FusedAttnBackend = { - "AOTriton": NVTE_Fused_Attn_Backend.NVTE_AOTriton, - "CK": NVTE_Fused_Attn_Backend.NVTE_CK, - "No_Backend": NVTE_Fused_Attn_Backend.NVTE_No_Backend, - } + +class FusedAttnBackend(IntEnum): + """Fused attention sub-backends. + + This is the canonical fused-attention backend enum for + ``transformer_engine.pytorch``. It mirrors the backend + ``transformer_engine_torch.NVTE_Fused_Attn_Backend`` (pybind11) enum + value-for-value, and instances of the two enums compare equal when they + share the same integer value. Unlike the pybind enum, a plain-python + ``IntEnum`` is traceable by ``torch.compile``: comparisons constant-fold + cleanly and instances safely cross the ``assume_constant_result`` boundary + in ``get_attention_backend``. Lookup by name (``FusedAttnBackend["FP8"]``) + works the same way as with the dict this used to be. + """ + + No_Backend = int(NVTE_Fused_Attn_Backend.NVTE_No_Backend) + if not IS_HIP_EXTENSION: + # The ROCm fork retains the (implementation-less) max512 enumerator on the + # CUDA C++ side for binding/consumer compile compatibility; mirror it here + # so the import-time sync assertion below matches the C++ enum. + F16_max512_seqlen = int(NVTE_Fused_Attn_Backend.NVTE_F16_max512_seqlen) + F16_arbitrary_seqlen = int(NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen) + FP8 = int(NVTE_Fused_Attn_Backend.NVTE_FP8) + else: + AOTriton = int(NVTE_Fused_Attn_Backend.NVTE_AOTriton) + CK = int(NVTE_Fused_Attn_Backend.NVTE_CK) + + @classmethod + def cast( + cls, backend: "Union[FusedAttnBackend, NVTE_Fused_Attn_Backend]" + ) -> "FusedAttnBackend": + """Normalize a backend value to the canonical ``FusedAttnBackend`` member. + + The pybind ``transformer_engine_torch.NVTE_Fused_Attn_Backend`` enum is + accepted as input for backward compatibility and mapped to the matching + ``FusedAttnBackend`` member. + """ + if isinstance(backend, cls): + return backend + return cls(int(backend)) + + def __eq__(self, other: object) -> bool: + # ``FusedAttnBackend`` is an ``IntEnum`` while ``NVTE_Fused_Attn_Backend`` + # is a pybind11 enum. Compare by integer value so the two enums stay + # equivalent regardless of the pybind11 version (the pybind ``__eq__`` + # handles the reverse order). + if isinstance(other, NVTE_Fused_Attn_Backend): + return int(self) == int(other) + return int.__eq__(self, other) + + def __ne__(self, other: object) -> bool: + result = self.__eq__(other) + if result is NotImplemented: + return result + return not result + + def __hash__(self) -> int: + return int.__hash__(self) + + +# Fail fast at import time if a new enumerator is added on the C++ side +# without being mirrored above. +assert {f"NVTE_{m.name}" for m in FusedAttnBackend} == set(NVTE_Fused_Attn_Backend.__members__), ( + "FusedAttnBackend in python is out of sync with" + " transformer_engine_torch.NVTE_Fused_Attn_Backend defined on the C++ side." + " Please make sure TE C++ and python are in sync." +) BACKEND_F16m512_FP8_THREADS_PER_CTA = 128 BACKEND_F16arb_ELTS_PER_THREADS = 16 @@ -135,7 +190,7 @@ def fused_attn_fwd( k: torch.Tensor, v: torch.Tensor, fake_dtype: torch.dtype, - fused_attention_backend: tex.NVTE_Fused_Attn_Backend, + fused_attention_backend: FusedAttnBackend, attn_bias: torch.Tensor = None, cu_seqlens_q_padded: torch.Tensor = None, cu_seqlens_kv_padded: torch.Tensor = None, @@ -187,7 +242,7 @@ def fused_attn_fwd( fake_dtype : DType data type of Q, K and V - in case of high precision, fake dtype in case of FP8; in torch.dtype - fused_attention_backend : tex.NVTE_Fused_Attn_Backend + fused_attention_backend : FusedAttnBackend please see FusedAttention module for details on supported backends. attn_bias : torch.Tensor, default = None input tensor Bias when attn_bias_type is "pre_scale_bias" or "post_scale_bias"; @@ -299,6 +354,8 @@ def fused_attn_fwd( f"attn_bias.dtype={attn_bias.dtype} but q.dtype={q.dtype}." ) + # Accept the pybind enum for backward compatibility. + fused_attention_backend = FusedAttnBackend.cast(fused_attention_backend) if fused_attention_backend == FusedAttnBackend["No_Backend"]: raise ValueError( "Fused attention does not support this input combination:" @@ -420,7 +477,7 @@ def fused_attn_bwd( d_o: torch.Tensor, fake_dtype: torch.dtype, aux_ctx_tensors: List[torch.Tensor], - fused_attention_backend: tex.NVTE_Fused_Attn_Backend, + fused_attention_backend: FusedAttnBackend, cu_seqlens_q_padded: torch.Tensor = None, cu_seqlens_kv_padded: torch.Tensor = None, s_quantizer: Quantizer = None, @@ -476,7 +533,7 @@ def fused_attn_bwd( aux_ctx_tensors : List[torch.Tensor] auxiliary output tensors of the forward pass when its is_training is True, e.g. aux_ctx_tensors = [S, Max, rng_state] - fused_attention_backend : tex.NVTE_Fused_Attn_Backend + fused_attention_backend : FusedAttnBackend please see FusedAttention module for details on supported backends. cu_seqlens_q_padded : torch.Tensor, default = None cumulative sequence offsets for Q; shape [batch_size + 1] @@ -561,6 +618,8 @@ def fused_attn_bwd( d = q.size(-1) attn_scale = 1.0 / math.sqrt(d) + # Accept the pybind enum for backward compatibility. + fused_attention_backend = FusedAttnBackend.cast(fused_attention_backend) if fused_attention_backend == FusedAttnBackend["No_Backend"]: raise ValueError( "Fused attention backward does not support this input combination:" diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 662ae9aa40..7203ccb353 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -6,7 +6,7 @@ """Python interface for GEMM extensions""" -from typing import Iterable, Optional, Tuple, Union, List +from typing import Iterable, Literal, Optional, Tuple, Union, List import os import functools import torch @@ -18,14 +18,20 @@ from ..utils import get_device_compute_capability from ..utils import cast_if_needed -from ..quantized_tensor import Quantizer +from ..quantized_tensor import QuantizedTensorStorage, Quantizer +from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer +from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer +from ..tensor.mxfp8_tensor import MXFP8Quantizer +from ..tensor.nvfp4_tensor import NVFP4Quantizer from ..tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage -from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage +from ..tensor.storage.float8_tensor_storage import Float8TensorStorage from ..tensor.storage.grouped_tensor_storage import GroupedTensorStorage +from ..tensor.storage.hybrid_tensor_storage import HybridQuantizedTensorStorage +from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from ..tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage from ..tensor.utils import is_custom from ..custom_recipes.gemm import custom_gemm -from ...debug.pytorch.debug_quantization import DebugQuantizer +from ...debug.pytorch.debug_quantization import DebugQuantizedTensor, DebugQuantizer _FP4_USE_TUNED_GEMM = int(os.environ.get("NVTE_FP4_USE_TUNED_GEMM", "1")) _FP4_LOG_SHAPES = int(os.environ.get("NVTE_FP4_LOG_GEMM_SHAPES", "0")) @@ -312,28 +318,108 @@ def _nvfp4_row_scaled_gemm_inputs( B: NVFP4TensorStorage, *, transa: bool, -) -> Tuple[NVFP4TensorStorage, NVFP4TensorStorage, torch.Tensor]: - """Return GEMM aliases and FP32 output scales for row-scaled NVFP4.""" + transb: bool, +) -> Tuple[NVFP4TensorStorage, NVFP4TensorStorage, torch.Tensor, torch.Tensor]: + """Return per-tensor GEMM aliases and row/column FP32 output scales.""" A_metadata = A.get_metadata() - weight_amax = A._amax_rowwise if transa else A._amax_columnwise - assert weight_amax is not None and weight_amax.numel() == 1 - A_metadata["amax_rowwise" if transa else "amax_columnwise"] = weight_amax.new_ones(1) + a_amax_key = "amax_rowwise" if transa else "amax_columnwise" + output_col_scales = A_metadata[a_amax_key] + assert output_col_scales is not None + A_metadata[a_amax_key] = output_col_scales.new_ones(1) A_metadata["row_scaled_nvfp4"] = False B_metadata = B.get_metadata() - rhs_rowwise_amax = B._amax_rowwise - assert rhs_rowwise_amax is not None - B_metadata["amax_rowwise"] = rhs_rowwise_amax.new_ones(1) + b_amax_key = "amax_columnwise" if transb else "amax_rowwise" + output_row_scales = B_metadata[b_amax_key] + assert output_row_scales is not None + B_metadata[b_amax_key] = output_row_scales.new_ones(1) B_metadata["row_scaled_nvfp4"] = False - assert rhs_rowwise_amax.dtype == torch.float32 and weight_amax.dtype == torch.float32 + assert output_row_scales.dtype == torch.float32 and output_col_scales.dtype == torch.float32 return ( NVFP4TensorStorage(**A_metadata), NVFP4TensorStorage(**B_metadata), - (rhs_rowwise_amax * weight_amax).view(-1, 1), + output_row_scales.view(-1, 1), + output_col_scales.view(1, -1), ) +_NATIVE_GEMM_INPUT_STORAGES = ( + Float8TensorStorage, + MXFP8TensorStorage, + Float8BlockwiseQTensorStorage, + NVFP4TensorStorage, +) + + +def _unwrap_tensor( + tensor: Union[torch.Tensor, QuantizedTensorStorage], + usage: Literal["rowwise", "columnwise"], +) -> Union[torch.Tensor, QuantizedTensorStorage]: + """Prepare a tensor for native or custom GEMM dispatch.""" + if usage not in ("rowwise", "columnwise"): + raise ValueError(f"Unsupported GEMM tensor usage ({usage})") + + # Hybrid and debug wrappers may omit a representation that is not needed + # by their configured GEMMs. Fail here if a caller requests that missing + # direction instead of passing ``None`` deeper into GEMM preparation. + if tensor is None: + raise RuntimeError( + f"GEMM requested the {usage} representation, but it is unavailable. " + f"Ensure {usage}_usage is enabled and the representation has not " + "been dropped by update_usage()." + ) + + # Plain PyTorch tensor + if not isinstance(tensor, QuantizedTensorStorage): + return tensor + + # Select the requested representation from a debug wrapper, then process + # the wrapped tensor normally (it may itself be hybrid or TE-native). + if isinstance(tensor, DebugQuantizedTensor): + return _unwrap_tensor(tensor.get_tensor(usage == "columnwise"), usage) + + # Select the direction of a hybrid tensor, then process its sub-storage. + if isinstance(tensor, HybridQuantizedTensorStorage): + sub_storage = ( + tensor.rowwise_sub_storage if usage == "rowwise" else tensor.columnwise_sub_storage + ) + return _unwrap_tensor(sub_storage, usage) + + # Preserve custom tensors for custom_gemm dispatch. + if is_custom(tensor): + return tensor + + # Quantized tensor formats with native GEMM support + if isinstance(tensor, _NATIVE_GEMM_INPUT_STORAGES): + return tensor + + # Fall back to high-precision GEMM for other quantized tensor formats. + return tensor.dequantize() + + +_NATIVE_GEMM_OUTPUT_QUANTIZERS = ( + Float8Quantizer, + Float8CurrentScalingQuantizer, + MXFP8Quantizer, + Float8BlockQuantizer, + NVFP4Quantizer, +) + + +def _validate_native_gemm_output_quantizer(quantization_params): + """Validate that the native C++ GEMM path can convert an output quantizer.""" + if quantization_params is not None and not isinstance( + quantization_params, _NATIVE_GEMM_OUTPUT_QUANTIZERS + ): + raise NotImplementedError( + f"{type(quantization_params).__name__} is not supported as a native " + "GEMM output quantizer. " + "Return a TE-native quantizer for output/grad_input roles or disable " + "quantized GEMM output for this boundary." + ) + + def general_gemm( A: torch.Tensor, B: torch.Tensor, @@ -360,6 +446,14 @@ def general_gemm( transa = layout[0] == "T" transb = layout[1] == "T" + debug_quantizer = None + if isinstance(quantization_params, DebugQuantizer): + debug_quantizer = quantization_params + quantization_params = quantization_params.parent_quantizer + + A = _unwrap_tensor(A, "rowwise" if transa else "columnwise") + B = _unwrap_tensor(B, "columnwise" if transb else "rowwise") + alpha = validate_gemm_scale(alpha, True) beta = validate_gemm_scale(beta, accumulate) @@ -417,12 +511,7 @@ def general_gemm( grad, ) - debug_quantizer = None - if isinstance(quantization_params, DebugQuantizer): - debug_quantizer = quantization_params - quantization_params = quantization_params.parent_quantizer - A = A.get_tensor(not transa) - B = B.get_tensor(transb) + _validate_native_gemm_output_quantizer(quantization_params) # Use bfloat16 as default bias_dtype bias_dtype = TE_DType[torch.bfloat16 if bias is None else bias.dtype] @@ -508,16 +597,7 @@ def general_gemm( elif not _is_nvfp4_row_scaled_tensor(A) and not _is_nvfp4_row_scaled_tensor(B): out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*args, **kwargs) else: - if _is_nvfp4_row_scaled_tensor(A): - raise NotImplementedError("Row-scaled NVFP4 GEMM does not support row-scaled A.") - assert layout[1] == "N", "Row-scaled NVFP4 GEMM currently supports N-layout B only." - if grad: - raise RuntimeError( - "Row-scaled NVFP4 GEMM currently supports fprop only. " - "Backward NVFP4 gradient quantizers should use scalar global amax." - ) assert not gelu, "Row-scaled NVFP4 GEMM currently does not support fused GELU." - assert not accumulate, "Row-scaled NVFP4 GEMM currently does not support accumulation." assert ( quantization_params is None ), "Row-scaled NVFP4 GEMM currently does not support output quantization." @@ -532,9 +612,14 @@ def general_gemm( assert isinstance( A, NVFP4TensorStorage ), "Row-scaled NVFP4 GEMM currently requires NVFP4 A." - # cuBLAS folds NVFP4 global amax values into GEMM alpha. Keep the row-scaled - # recipe's global scales out of alpha and apply them in FP32 below. - gemm_A, gemm_B, rowwise_global_scales = _nvfp4_row_scaled_gemm_inputs(A, B, transa=transa) + assert isinstance( + B, NVFP4TensorStorage + ), "Row-scaled NVFP4 GEMM currently requires NVFP4 B." + # Reuse the per-tensor GEMM and apply selected row/column global scales + # to the FP32 output. This extends #2931 without a dedicated GEMM kernel. + gemm_A, gemm_B, output_row_scales, output_col_scales = _nvfp4_row_scaled_gemm_inputs( + A, B, transa=transa, transb=transb + ) requested_out, requested_out_dtype = out, out_dtype fp32_out = ( @@ -549,18 +634,36 @@ def general_gemm( gemm_args[5] = None # quantization_params gemm_args[6] = TE_DType[torch.float32] # out_dtype gemm_args[7] = None # bias - out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*gemm_args, **kwargs) + gemm_args[14] = False # accumulate after applying the outer scales + gemm_kwargs = dict(kwargs) + gemm_kwargs["beta"] = 0.0 + out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*gemm_args, **gemm_kwargs) out_2d = out.reshape(-1, out.shape[-1]) - assert rowwise_global_scales.dtype == torch.float32 and out.dtype == torch.float32 - assert rowwise_global_scales.numel() == out_2d.shape[0] - - out_2d.mul_(rowwise_global_scales) + assert output_row_scales.numel() in (1, out_2d.shape[0]) + assert output_col_scales.numel() in (1, out_2d.shape[1]) + assert out.dtype == torch.float32 + # When one side is a scalar global amax (e.g. fprop weight), fold both + # scales into a single factor before the multiply. This reproduces + # #2931's fused `out * (row_amax * col_amax)` arithmetic bit-for-bit; + # only the true bilateral case (both per-row and per-col) needs the + # two-step outer-product scaling. + if output_col_scales.numel() == 1: + out_2d.mul_(output_row_scales * output_col_scales) + elif output_row_scales.numel() == 1: + out_2d.mul_(output_col_scales * output_row_scales) + else: + out_2d.mul_(output_row_scales) + out_2d.mul_(output_col_scales) if bias is not None: + assert not grad, "Row-scaled NVFP4 backward does not support fused bias gradient." out_2d.add_(bias.to(dtype=torch.float32)) if requested_out is not None: - requested_out.copy_(out.to(dtype=requested_out.dtype)) + if accumulate: + requested_out.add_(out.to(dtype=requested_out.dtype)) + else: + requested_out.copy_(out.to(dtype=requested_out.dtype)) out = requested_out elif requested_out_dtype is not None and requested_out_dtype != torch.float32: out = out.to(dtype=requested_out_dtype) @@ -599,6 +702,9 @@ def general_grouped_gemm( transa = layout[0] == "T" transb = layout[1] == "T" + A = [_unwrap_tensor(a, "rowwise" if transa else "columnwise") for a in A] + B = [_unwrap_tensor(b, "columnwise" if transb else "rowwise") for b in B] + empty_tensor = _empty_tensor() empty_tensors = [empty_tensor] * num_gemms @@ -621,9 +727,9 @@ def general_grouped_gemm( else: bias_dtype = TE_DType[torch.bfloat16] - if any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in A): - raise NotImplementedError("Row-scaled NVFP4 grouped GEMM does not support row-scaled A.") - if any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in B): + if any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in A) or any( + _is_nvfp4_row_scaled_tensor(tensor) for tensor in B + ): assert D_dtype is None, "Row-scaled NVFP4 grouped GEMM currently does not support D_dtype." if single_output: assert ( @@ -736,6 +842,30 @@ def _get_fp32_zeros_tensor(num_tensors: int, device: torch.device) -> torch.Tens return torch.zeros(num_tensors, dtype=torch.float32, device=device) +@functools.lru_cache(maxsize=None) +def _get_grouped_gemm_setup_workspace(device: int, num_tensors: int) -> torch.Tensor: + """Persistent setup workspace (per-group pointer/dim arrays) for grouped-tensor GEMM.""" + return torch.empty( + get_grouped_gemm_setup_workspace_size(num_tensors), + dtype=torch.uint8, + device=device, + ) + + +@functools.lru_cache(maxsize=None) +def _get_grouped_cublas_workspace(device: int, layout: str) -> torch.Tensor: + """Persistent cuBLAS workspace for the grouped-tensor GEMM path, one per GEMM layout. + + Grouped cuBlasLt GEMM kernels in cuBLAS versions <= 13.7 leave behind stale descriptors in the + workspace that cause back-to-back GEMM kernels to crash/deadlock on 2nd CUDA-graph replay. As a + workaround, we allocate a different workspace for each GEMM layout (TN, NN, NT) to avoid + contamination between subsequent GEMM calls (when there is no other graph node between GEMM + kernels). + """ + assert layout in ("TN", "NN", "NT"), f"unexpected grouped GEMM layout {layout}" + return torch.empty(get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device) + + def general_grouped_gemm_for_grouped_tensor( A, B, @@ -775,6 +905,20 @@ def general_grouped_gemm_for_grouped_tensor( if isinstance(out, GroupedTensorStorage) and out.row_scaled_nvfp4: raise NotImplementedError("Row-scaled NVFP4 GroupedTensor GEMM is not supported yet.") + def _is_fp8_blockwise(operand) -> bool: + if isinstance(operand, (list, tuple)): + return any(isinstance(t, Float8BlockwiseQTensorStorage) for t in operand) + if isinstance(operand, GroupedTensorStorage): + return isinstance(operand.quantizer, Float8BlockQuantizer) + return False + + if _is_fp8_blockwise(A) or _is_fp8_blockwise(B): + # The fused grouped FP8 block-scaling GEMM only supports split accumulation, + # so force it on and intentionally override any caller-supplied value. This + # matches the Float8BlockScaling recipe, which fixes use_split_accumulator=True + # for all of fprop/dgrad/wgrad, so no user-configurable setting is discarded. + use_split_accumulator = True + if is_discrete_out: # wgrad case. grouped_gemm_impl = tex.te_general_grouped_gemm_for_discrete_out @@ -814,16 +958,12 @@ def general_grouped_gemm_for_grouped_tensor( if not alpha.is_cuda or not beta.is_cuda: raise ValueError("alpha and beta must be CUDA tensors.") - workspace_setup = torch.empty( - get_grouped_gemm_setup_workspace_size(num_tensors), - dtype=torch.uint8, - device=device, - ) - workspace_cublas = torch.empty( - get_cublas_workspace_size_bytes(), - dtype=torch.uint8, - device=device, - ) + workspace_setup = _get_grouped_gemm_setup_workspace(device.index, num_tensors) + # Each grouped-GEMM layout gets its own persistent cuBLAS workspace: two grouped + # GEMMs sharing one workspace can deadlock under CUDA-graph replay (see + # _get_grouped_cublas_workspace). wgrad (NT) is the case seen in TE; fprop (TN) and + # dgrad (NN) have also been reported to conflict, so all three layouts are isolated. + workspace_cublas = _get_grouped_cublas_workspace(device.index, layout) sm_count = get_sm_count() sm_count = sm_count - int(os.getenv("NVTE_EXT_MARGIN_SM", str(sm_count))) diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index cef96c6efd..16a9965d51 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -320,6 +320,8 @@ class Float8BlockQuantizer : public Quantizer { class MXFP8Quantizer : public Quantizer { public: + bool with_2d_quantization = false; + explicit MXFP8Quantizer(const py::handle& quantizer); NVTEScalingMode get_scaling_mode() const override { return NVTE_MXFP8_1D_SCALING; } diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index a1263de7f4..2693a7ed86 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -292,6 +292,31 @@ py::object clamped_swiglu(const at::Tensor &input, py::handle quantizer, float l py::object clamped_dswiglu(const at::Tensor &grad, const at::Tensor &input, py::handle quantizer, float limit, float alpha, float glu_linear_offset); + +/* Scaled activation */ +py::object scaled_swiglu(const at::Tensor &input, const at::Tensor &act_scales, + py::handle quantizer, int64_t glu_interleave_size); + +py::object scaled_clamped_swiglu(const at::Tensor &input, const at::Tensor &act_scales, + py::handle quantizer, float limit, float alpha, + float glu_linear_offset, int64_t glu_interleave_size); + +py::object scaled_srelu(const at::Tensor &input, const at::Tensor &act_scales, + py::handle quantizer); + +py::tuple scaled_dswiglu(const at::Tensor &grad, const at::Tensor &input, + const at::Tensor &act_scales, py::handle quantizer, + int64_t glu_interleave_size, bool compute_scale_grad); + +py::tuple scaled_clamped_dswiglu(const at::Tensor &grad, const at::Tensor &input, + const at::Tensor &act_scales, py::handle quantizer, float limit, + float alpha, float glu_linear_offset, int64_t glu_interleave_size, + bool compute_scale_grad); + +py::tuple scaled_dsrelu(const at::Tensor &grad, const at::Tensor &input, + const at::Tensor &act_scales, py::handle quantizer, + bool compute_scale_grad); + /*************************************************************************************************** * LayerNorm **************************************************************************************************/ @@ -352,7 +377,7 @@ py::object dequantize(const py::handle &input, DType otype); py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, std::optional first_dims, std::optional last_dims, std::optional tensor_offsets, - std::optional noop_flag); + std::optional noop_flag, const py::object &output); py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, @@ -369,6 +394,11 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, std::optional last_dims, std::optional tensor_offsets); +py::object group_requantize_inplace(py::handle grouped_x, py::handle quantizer, + const size_t num_tensors, std::optional first_dims, + DType otype, std::optional tensor_offsets, + bool return_dequantized); + std::vector multi_tensor_quantize(const std::vector &tensor_list, std::vector quantizer_list); @@ -722,7 +752,7 @@ void register_ep_bindings(pybind11::module_ &m); void ep_initialize(uintptr_t comm_ptr, const std::string &group_name, int64_t num_experts, int64_t max_tokens_per_rank, int64_t max_recv_tokens_per_rank, int64_t hidden_dim, int64_t max_num_sms, pybind11::object max_token_dtype, - bool zero_copy); + bool zero_copy, int64_t num_topk, bool drop_on_overflow); void ep_finalize(); @@ -732,18 +762,23 @@ bool ep_get_zero_copy(); // Returns the handle_mem byte size for the given layer config. int64_t ep_handle_mem_size(int64_t top_k, int64_t dispatch_output_per_expert_alignment); -void ep_prepare(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor token_counts, int64_t top_k, - int64_t dispatch_output_per_expert_alignment); +void ep_prepare(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens_per_expert, + int64_t top_k, int64_t dispatch_output_per_expert_alignment, + at::Tensor total_recv_tokens); void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens, - at::Tensor topk_weights, at::Tensor recv_tokens, at::Tensor recv_topk_weights); + at::Tensor topk_weights, at::Tensor recv_tokens, at::Tensor recv_topk_weights, + std::optional tokens_scale_inv = std::nullopt, + std::optional recv_scale_inv = std::nullopt); void ep_combine(at::Tensor handle_mem, at::Tensor expert_out, at::Tensor result); void ep_dispatch_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor g_recv_topk_weights, at::Tensor grad_tokens, at::Tensor grad_topk_weights); -void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expert_out); +void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expert_out, + std::optional grad_scale_inv = std::nullopt, + std::optional grad_expert_out_scale_inv = std::nullopt); // Registers the EP pybind functions on `m`. Defined under NVTE_WITH_NCCL_EP. void register_ep_bindings(pybind11::module_ &m); diff --git a/transformer_engine/pytorch/csrc/extensions/activation.cpp b/transformer_engine/pytorch/csrc/extensions/activation.cpp index 3a04a95150..d116e8bb10 100644 --- a/transformer_engine/pytorch/csrc/extensions/activation.cpp +++ b/transformer_engine/pytorch/csrc/extensions/activation.cpp @@ -356,5 +356,150 @@ py::object clamped_dswiglu(const at::Tensor& grad, const at::Tensor& input, py:: glu_linear_offset); } +/* Scaled activation helpers (activation + per-row scale via nvte_scaled_*). */ + +template +at::Tensor scaled_activation_compute(const at::Tensor& input, const at::Tensor& act_scales, + int shape_divisor, Args&&... args) { + init_extension(); + NVTE_CHECK(input.dim() >= 1, "scaled activation input must have at least 1 dimension"); + NVTE_CHECK(shape_divisor > 0 && input.size(-1) % shape_divisor == 0, + "scaled activation input width is not compatible with activation"); + + auto input_tensor = input.contiguous(); + auto scales_tensor = act_scales.contiguous().reshape({-1}); + const int64_t rows = input_tensor.numel() / input_tensor.size(-1); + NVTE_CHECK(scales_tensor.numel() == rows, "scaled activation expects one scale per input row"); + + std::vector output_sizes(input_tensor.sizes().begin(), input_tensor.sizes().end()); + output_sizes.back() /= shape_divisor; + auto output = at::empty(output_sizes, input_tensor.options()); + + const TensorWrapper& input_nvte = makeTransformerEngineTensor(input_tensor); + const TensorWrapper& scales_nvte = makeTransformerEngineTensor(scales_tensor); + const TensorWrapper& output_nvte = makeTransformerEngineTensor(output); + + auto stream = at::cuda::getCurrentCUDAStream(); + NVTE_SCOPED_GIL_RELEASE({ + act_func(input_nvte.data(), scales_nvte.data(), output_nvte.data(), std::forward(args)..., + stream); + }); + return output; +} + +template +std::tuple scaled_dactivation_compute(const at::Tensor& grad, + const at::Tensor& input, + const at::Tensor& act_scales, + bool compute_scale_grad, + Args&&... args) { + init_extension(); + NVTE_CHECK(input.dim() >= 1 && grad.dim() >= 1, + "scaled dactivation input and grad must have at least 1 dimension"); + + auto grad_tensor = grad.contiguous(); + auto input_tensor = input.contiguous(); + auto scales_tensor = act_scales.contiguous(); + const int64_t rows = input_tensor.numel() / input_tensor.size(-1); + NVTE_CHECK(scales_tensor.numel() == rows, "scaled dactivation expects one scale per input row"); + + auto scales_flat = scales_tensor.reshape({-1}); + auto grad_input = at::empty_like(input_tensor); + auto grad_scales = compute_scale_grad ? at::empty_like(scales_tensor) : at::Tensor(); + auto grad_scales_flat = compute_scale_grad ? grad_scales.reshape({-1}) : at::Tensor(); + + const TensorWrapper& grad_nvte = makeTransformerEngineTensor(grad_tensor); + const TensorWrapper& input_nvte = makeTransformerEngineTensor(input_tensor); + const TensorWrapper& scales_nvte = makeTransformerEngineTensor(scales_flat); + const TensorWrapper& grad_input_nvte = makeTransformerEngineTensor(grad_input); + std::optional grad_scales_nvte; + if (compute_scale_grad) { + grad_scales_nvte.emplace(makeTransformerEngineTensor(grad_scales_flat)); + } + + auto stream = at::cuda::getCurrentCUDAStream(); + NVTE_SCOPED_GIL_RELEASE({ + dact_func(grad_nvte.data(), input_nvte.data(), scales_nvte.data(), grad_input_nvte.data(), + compute_scale_grad ? grad_scales_nvte->data() : nullptr, std::forward(args)..., + stream); + }); + return {grad_input, grad_scales}; +} + +py::object maybe_quantize(const at::Tensor& tensor, py::handle quantizer) { + if (quantizer.is_none()) { + return py::cast(tensor); + } + auto quantizer_cpp = convert_quantizer(quantizer); + const TensorWrapper& tensor_nvte = makeTransformerEngineTensor(tensor); + const auto shape_te = tensor_nvte.shape(); + const std::vector shape(shape_te.data, shape_te.data + shape_te.ndim); + auto fake_dtype = GetTransformerEngineDType(tensor.scalar_type()); + auto [out_nvte, out_py] = quantizer_cpp->create_tensor(shape, fake_dtype); + quantizer_cpp->quantize(tensor_nvte, out_nvte); + return out_py; +} + +template +py::object scaled_activation_helper(const at::Tensor& input, const at::Tensor& act_scales, + py::handle quantizer, int shape_divisor, Args&&... args) { + auto output = scaled_activation_compute(input, act_scales, shape_divisor, + std::forward(args)...); + return maybe_quantize(output, quantizer); +} + +template +py::tuple scaled_dactivation_helper(const at::Tensor& grad, const at::Tensor& input, + const at::Tensor& act_scales, py::handle quantizer, + bool compute_scale_grad, Args&&... args) { + auto [grad_input, grad_scales] = scaled_dactivation_compute( + grad, input, act_scales, compute_scale_grad, std::forward(args)...); + return py::make_tuple(maybe_quantize(grad_input, quantizer), + compute_scale_grad ? py::cast(grad_scales) : py::none()); +} + +py::object scaled_swiglu(const at::Tensor& input, const at::Tensor& act_scales, + py::handle quantizer, int64_t glu_interleave_size) { + return scaled_activation_helper(input, act_scales, quantizer, + /*shape_divisor=*/2, glu_interleave_size); +} + +py::object scaled_clamped_swiglu(const at::Tensor& input, const at::Tensor& act_scales, + py::handle quantizer, float limit, float alpha, + float glu_linear_offset, int64_t glu_interleave_size) { + return scaled_activation_helper( + input, act_scales, quantizer, /*shape_divisor=*/2, limit, alpha, glu_linear_offset, + glu_interleave_size); +} + +py::object scaled_srelu(const at::Tensor& input, const at::Tensor& act_scales, + py::handle quantizer) { + return scaled_activation_helper(input, act_scales, quantizer, + /*shape_divisor=*/1); +} + +py::tuple scaled_dswiglu(const at::Tensor& grad, const at::Tensor& input, + const at::Tensor& act_scales, py::handle quantizer, + int64_t glu_interleave_size, bool compute_scale_grad) { + return scaled_dactivation_helper(grad, input, act_scales, quantizer, + compute_scale_grad, glu_interleave_size); +} + +py::tuple scaled_clamped_dswiglu(const at::Tensor& grad, const at::Tensor& input, + const at::Tensor& act_scales, py::handle quantizer, float limit, + float alpha, float glu_linear_offset, int64_t glu_interleave_size, + bool compute_scale_grad) { + return scaled_dactivation_helper( + grad, input, act_scales, quantizer, compute_scale_grad, limit, alpha, glu_linear_offset, + glu_interleave_size); +} + +py::tuple scaled_dsrelu(const at::Tensor& grad, const at::Tensor& input, + const at::Tensor& act_scales, py::handle quantizer, + bool compute_scale_grad) { + return scaled_dactivation_helper(grad, input, act_scales, quantizer, + compute_scale_grad); +} + } // namespace pytorch } // namespace transformer_engine diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index e21ac6940a..ae8dcd9499 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -10,6 +10,7 @@ #include "common.h" #include "pybind.h" +// ROCm: at::cuda::CUDAGuard is unavailable on older torch; masquerade the HIP guard. #include #if USE_ROCM && TORCH_VERSION_MINOR < 11 using TECUDAGuard = at::hip::HIPGuardMasqueradingAsCUDA; @@ -17,35 +18,6 @@ using TECUDAGuard = at::hip::HIPGuardMasqueradingAsCUDA; using TECUDAGuard = at::cuda::CUDAGuard; #endif -namespace { - -constexpr int block_size = 512; - -// fast zero-fills of tensors -void mha_fill(const transformer_engine::TensorWrapper &self, const at::Tensor &start_index) { - std::vector shape = transformer_engine::pytorch::convertShape(self.shape()); - - auto max_tokens = shape[0]; - auto fcd_size = 1; - for (size_t i = 1; i <= shape.size(); i++) { - fcd_size *= shape[i]; - } - - NVTE_CHECK(fcd_size % block_size == 0, "input size not aligned to block size"); - - size_t element_size_bits = transformer_engine::pytorch::typeToNumBits(self.dtype()); - int32_t start_row = start_index.data_ptr()[0]; - void *base_ptr = static_cast(self.get_rowwise_data().data_ptr) + - static_cast(start_row) * fcd_size * element_size_bits / 8; - size_t num_rows_to_zero = max_tokens - start_row; - size_t total_bytes = num_rows_to_zero * fcd_size * element_size_bits / 8; - - NVTE_SCOPED_GIL_RELEASE( - { nvte_memset(base_ptr, 0, total_bytes, at::cuda::getCurrentCUDAStream()); }); -} - -} // namespace - namespace transformer_engine::pytorch { // get the fused attention backend @@ -159,7 +131,6 @@ std::vector fused_attn_fwd( auto o_shape = std::vector{o_shape_tmp.begin(), o_shape_tmp.end()}; NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); AttentionShape o_parsed(q_format, o_shape_tmp.data()); - size_t h = o_parsed.h(), d = o_parsed.d(); o_parsed.to_format(o_format, o_shape.data()); const DType fake_dtype_te = GetTransformerEngineDType(fake_dtype); auto [te_O, py_O, o_amax_buf] = @@ -173,11 +144,8 @@ std::vector fused_attn_fwd( if (qkv_type == DType::kFloat8E4M3 || qkv_type == DType::kFloat8E5M2) { // FP8 if (set_zero && (o_format == NVTE_QKV_Format::NVTE_THD)) { - if ((h * d) % block_size == 0) { - mha_fill(te_O, cu_seqlens_q.index({torch::indexing::Slice(-1, torch::indexing::None)})); - } else { - te_O.zero_(at::cuda::getCurrentCUDAStream()); - } + // Initialize both the output data and its amax metadata. + te_O.zero_(at::cuda::getCurrentCUDAStream()); } } else if (qkv_type == DType::kBFloat16 || qkv_type == DType::kFloat16) { if (o_format == NVTE_QKV_Format::NVTE_THD) { @@ -374,13 +342,11 @@ std::vector fused_attn_bwd( NVTE_QKV_Format dq_format = nvte_get_q_format(dqkv_layout); NVTE_QKV_Format dkv_format = nvte_get_kv_format(dqkv_layout); AttentionShape q_parsed(q_format, q_shape.data()); - size_t h_q = q_parsed.h(), d_qk = q_parsed.d(); + size_t h_q = q_parsed.h(); q_parsed.to_format(dq_format, dQ_shape.data()); AttentionShape k_parsed(kv_format, k_shape.data()); - size_t h_kv = k_parsed.h(); k_parsed.to_format(dkv_format, dK_shape.data()); AttentionShape v_parsed(kv_format, v_shape.data()); - size_t d_v = v_parsed.d(); v_parsed.to_format(dkv_format, dV_shape.data()); at::Tensor dQ, dK, dV, dQKV, dKV; // FP16/BF16: dqkv_fake_dtype = kFloat16/kBFloat16, dQ/dK/dV.dtype = torch.float16/torch.bfloat16 @@ -477,22 +443,18 @@ std::vector fused_attn_bwd( if (detail::IsFloat8Quantizers(dqkv_quantizer.ptr())) { // FP8 if (set_zero) { + // dQ/dK/dV may be strided views, so zero data through ATen and reset the shared amax + // separately. if (dq_format == NVTE_QKV_Format::NVTE_THD) { - if (((h_q * d_qk) % block_size == 0) && dQ.is_contiguous()) { - mha_fill(te_dQ, cu_seqlens_q.index({torch::indexing::Slice(-1, torch::indexing::None)})); - } else { - dQ.fill_(0); - } + dQ.fill_(0); } if (dkv_format == NVTE_QKV_Format::NVTE_THD) { - if (((h_kv * d_qk) % block_size == 0) && ((h_kv * d_v) % block_size == 0) && - dK.is_contiguous() && dV.is_contiguous()) { - mha_fill(te_dK, cu_seqlens_kv.index({torch::indexing::Slice(-1, torch::indexing::None)})); - mha_fill(te_dV, cu_seqlens_kv.index({torch::indexing::Slice(-1, torch::indexing::None)})); - } else { - dK.fill_(0); - dV.fill_(0); - } + dK.fill_(0); + dV.fill_(0); + } + if (dq_format == NVTE_QKV_Format::NVTE_THD || dkv_format == NVTE_QKV_Format::NVTE_THD) { + auto *fp8_quantizer = dynamic_cast(dQKV_quantizer.get()); + fp8_quantizer->amax.zero_(); } } } else if (dqkv_quantizer.is_none() || diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 2b4ec26576..adbc574f68 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -22,6 +22,7 @@ #include "../extensions.h" #include "common.h" #include "common/common.h" +#include "common/util/cuda_runtime.h" #include "common/util/system.h" #include "pybind.h" #include "transformer_engine/multi_tensor.h" @@ -281,7 +282,7 @@ void compute_grouped_fp8_current_scaling_amax_and_scale( py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, std::optional first_dims, std::optional last_dims, std::optional tensor_offsets, - std::optional noop_flag) { + std::optional noop_flag, const py::object &output) { using namespace transformer_engine::pytorch::detail; init_extension(); @@ -310,11 +311,34 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const GetTransformerEngineDType(tensor.scalar_type()), std::vector{static_cast(tensor.numel())}); - // Create output GroupedTensor. - auto [grouped_output_tensor_cpp, grouped_output_py] = quantizer_cpp->create_grouped_tensor( - num_tensors, logical_shape, GetTransformerEngineDType(tensor.scalar_type()), - py::reinterpret_borrow(quantizer), first_dims, last_dims, tensor_offsets, - logical_first_dim, logical_last_dim); + // Create a GroupedTensor or reuse an existing destination. Reusing the destination is + // required for weight caching and CUDA graph replay because captured GEMMs retain the + // original data and scale pointers. + py::object grouped_output_py; + auto grouped_output_tensor_cpp = [&]() -> GroupedTensorWrapper { + if (!output.is_none()) { + NVTE_CHECK(!first_dims.has_value() && !last_dims.has_value() && !tensor_offsets.has_value(), + "group_quantize: output reuse currently requires uniform tensor shapes."); + NVTE_CHECK(output.attr("num_tensors").cast() == num_tensors, + "group_quantize: output has a different number of tensors."); + NVTE_CHECK(output.attr("logical_shape").cast>() == logical_shape, + "group_quantize: output has an incompatible logical shape."); + py::object output_quantizer = output.attr("quantizer"); + NVTE_CHECK(!output_quantizer.is_none(), + "group_quantize: output must have quantized storage."); + NVTE_CHECK(output_quantizer.is(quantizer), + "group_quantize: output must have been created by the same quantizer."); + grouped_output_py = output; + return GroupedTensorFromPyTorchGroupedTensor(output); + } + + auto result = quantizer_cpp->create_grouped_tensor( + num_tensors, logical_shape, GetTransformerEngineDType(tensor.scalar_type()), + py::reinterpret_borrow(quantizer), first_dims, last_dims, tensor_offsets, + logical_first_dim, logical_last_dim); + grouped_output_py = std::move(result.second); + return std::move(result.first); + }(); // dispatch to scaling methods enum class GroupedQuantizationMode { @@ -378,10 +402,12 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const break; } case GroupedQuantizationMode::MXFP8_GROUPED_QUANTIZE: { + auto *mxfp8_quantizer_cpp = static_cast(quantizer_cpp.get()); QuantizationConfigWrapper quant_config_cpp; if (noop_flag_cpp.has_value()) { quant_config_cpp.set_noop_tensor(noop_flag_cpp->data()); } + quant_config_cpp.set_mxfp8_2d_quantization(mxfp8_quantizer_cpp->with_2d_quantization); NVTE_SCOPED_GIL_RELEASE({ nvte_group_quantize(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(), quant_config_cpp, at::cuda::getCurrentCUDAStream()); @@ -394,6 +420,9 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const QuantizationConfigWrapper quant_config_cpp; quant_config_cpp.set_force_pow_2_scales(fp8_block_quantizer_cpp->force_pow_2_scales); quant_config_cpp.set_amax_epsilon(fp8_block_quantizer_cpp->amax_epsilon); + if (noop_flag_cpp.has_value()) { + quant_config_cpp.set_noop_tensor(noop_flag_cpp->data()); + } NVTE_SCOPED_GIL_RELEASE({ nvte_group_quantize(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(), quant_config_cpp, at::cuda::getCurrentCUDAStream()); @@ -518,6 +547,15 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, auto quantizer_cpp = convert_quantizer(quantizer); + // MXFP8 carries no quantization knobs; FP8 block-scaling reads scale constraints + // off the quantizer, matching the forward group_quantize dispatch. + QuantizationConfigWrapper quant_config_cpp; + if (detail::IsFloat8BlockwiseQuantizers(quantizer.ptr())) { + auto *fp8_block_quantizer_cpp = static_cast(quantizer_cpp.get()); + quant_config_cpp.set_force_pow_2_scales(fp8_block_quantizer_cpp->force_pow_2_scales); + quant_config_cpp.set_amax_epsilon(fp8_block_quantizer_cpp->amax_epsilon); + } + auto grouped_input_tensor = GroupedTensorWrapper(num_tensors, logical_shape); grouped_input_tensor.set_rowwise_data(tensor.data_ptr(), GetTransformerEngineDType(tensor.scalar_type()), @@ -548,7 +586,8 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, auto stream = at::cuda::getCurrentCUDAStream(); NVTE_SCOPED_GIL_RELEASE({ nvte_group_quantize_dbias(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(), - grouped_dbias.data(), workspace_nvte.data(), stream); + grouped_dbias.data(), workspace_nvte.data(), quant_config_cpp, + stream); }); if (workspace_nvte.ndim() > 0 && workspace_nvte.numel() > 0) { at::Tensor workspace_torch = allocateSpace(workspace_nvte.shape(), workspace_nvte.dtype()); @@ -557,7 +596,8 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, } NVTE_SCOPED_GIL_RELEASE({ nvte_group_quantize_dbias(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(), - grouped_dbias.data(), workspace_nvte.data(), stream); + grouped_dbias.data(), workspace_nvte.data(), quant_config_cpp, + stream); }); return py::make_tuple(py::reinterpret_borrow(grouped_output_py), py::cast(std::move(dbias_torch))); @@ -670,6 +710,102 @@ py::object group_dequantize(const py::handle &input, transformer_engine::DType o return py::reinterpret_borrow(out_py); } +py::object group_requantize_inplace(py::handle grouped_x, py::handle quantizer, + const size_t num_tensors, std::optional first_dims, + DType otype, std::optional tensor_offsets, + bool return_dequantized) { + init_extension(); + + const bool has_rowwise = + !grouped_x.attr("rowwise_data").is_none() && !grouped_x.attr("scale_inv").is_none(); + const bool has_columnwise = !grouped_x.attr("columnwise_data").is_none() && + !grouped_x.attr("columnwise_scale_inv").is_none(); + const bool swizzled = grouped_x.attr("_with_gemm_swizzled_scales").cast(); + + NVTE_CHECK(has_rowwise, "Grouped input has no rowwise data and scales for the GEMM to consume."); + + // The tensor's own quantization must match what the op expects on every path: even a + // pass-through hands its data straight to the GEMM. This keeps the input's format rather than + // converting between formats. + const auto input_quantizer = grouped_x.attr("quantizer"); + NVTE_CHECK(!input_quantizer.is_none(), "Grouped input has no quantizer."); + NVTE_CHECK(Py_TYPE(input_quantizer.ptr()) == Py_TYPE(quantizer.ptr()), + "Grouped input and the op disagree on quantization format."); + NVTE_CHECK(input_quantizer.attr("dtype").cast() == quantizer.attr("dtype").cast(), + "Grouped input and the quantizer disagree on the FP8 dtype."); + + // The columnwise copy is only worth building when a wgrad GEMM will consume it. Read this + // before the usage is overridden below. + const bool need_columnwise = quantizer.attr("columnwise_usage").cast(); + + if (swizzled) { + // Already GEMM-ready. Nothing can be derived from here, since dequantization requires scales + // in compact format, so the input must already carry everything that will be consumed. No + // quantization kernel runs, which makes this path format-agnostic. + NVTE_CHECK(has_columnwise || !need_columnwise, + "Grouped input has swizzled scales but no columnwise data for the wgrad GEMM. It " + "cannot be rebuilt, because dequantization requires scales in compact format."); + NVTE_CHECK(!return_dequantized, + "Cannot return a dequantized tensor for an already-swizzled grouped input: " + "dequantization requires scales in compact format."); + return py::none(); + } + + NVTE_CHECK(!has_columnwise, + "Grouped input already has columnwise data but unswizzled scales; requantizing " + "from it is not supported."); + + // Everything below runs quantization kernels, so it is MXFP8-only. + NVTE_CHECK(detail::IsMXFP8Quantizers(quantizer.ptr()), + "Requantizing a grouped input is only supported for MXFP8."); + + const auto logical_shape = grouped_x.attr("logical_shape").cast(); + const auto total_tokens = logical_shape[0].cast(); + const auto hidden_dim = logical_shape[1].cast(); + // Each group's token count must be a multiple of 128 too, so that every group's scales start + // on a swizzle-tile boundary. Those counts live on the device (host reads would break CUDA + // graph capture), so that half is the caller's contract rather than an assertion. + NVTE_CHECK(total_tokens % 128 == 0 && hidden_dim % 128 == 0, + "Requantizing a grouped input requires dims that are multiples of 128, but got (", + total_tokens, ", ", hidden_dim, ")."); + + // Dequantize first: it reads the rowwise scales, which the swizzle below replaces. Left + // undefined when nothing consumes it, which skips the pass entirely. + at::Tensor dequantized; + if (need_columnwise || return_dequantized) { + dequantized = group_dequantize(grouped_x, otype) + .attr("rowwise_data") + .cast() + .view({static_cast(total_tokens), static_cast(hidden_dim)}); + } + + // Swizzle the rowwise scales before attaching any columnwise data: a rowwise-only swizzle + // resets columnwise_scale_inv to None, which would strand the columnwise data below with a + // null scale pointer. + grouped_swizzle_for_gemm(grouped_x, /*rowwise=*/true, /*columnwise=*/false); + // The swizzle hands back a 2D [num_tensors * padded_m, padded_k] scale buffer, but grouped + // tensors carry scales as a flat array indexed by element offsets (scale_inv_offsets), so + // per-group slicing breaks unless it is flattened back. + grouped_x.attr("scale_inv") = grouped_x.attr("scale_inv").attr("reshape")(-1); + + if (need_columnwise) { + // Rebuild the columnwise copy the wgrad GEMM needs. It cannot be derived from the rowwise + // data because the two directions scale along perpendicular axes. Quantizing rowwise as well + // would redo work we already have, so that direction is switched off; the caller sets + // optimize_for_gemm, which makes the kernel emit swizzled columnwise scales directly. + quantizer.attr("set_usage")(py::arg("rowwise") = false, py::arg("columnwise") = true); + auto columnwise = group_quantize(dequantized, quantizer, num_tensors, first_dims, std::nullopt, + tensor_offsets, std::nullopt, py::none()); + grouped_x.attr("columnwise_data") = columnwise.attr("columnwise_data"); + grouped_x.attr("columnwise_scale_inv") = columnwise.attr("columnwise_scale_inv"); + } + + if (return_dequantized) { + return py::cast(dequantized); + } + return py::none(); +} + namespace { void multi_tensor_quantize_impl(const std::vector &input_list, @@ -1039,8 +1175,6 @@ std::tuple, std::vector, bool> bulk_alloc const auto columnwise_usage = quantizer_cpp_list[0]->columnwise_usage; if (row_scaled_nvfp4) { NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 bulk allocation requires rowwise usage."); - NVTE_CHECK(!columnwise_usage, - "Row-scaled NVFP4 bulk allocation does not support columnwise usage."); } const auto scaling_mode = quantizer_cpp_list[0]->get_scaling_mode(); const auto fp4_dtype = quantizer_cpp_list[0]->dtype; @@ -1184,7 +1318,10 @@ std::tuple, std::vector, bool> bulk_alloc dtypes.insert(dtypes.end(), num_tensors, torch::kUInt8); alignments.insert(alignments.end(), num_tensors, 16); for (size_t i = 0; i < num_tensors; ++i) { - shapes.emplace_back(amax_shape(columnwise_data_shapes[i])); + // columnwise_data_shapes[i] is the transposed shape, so its leading dim is + // the original last dim (number of columns). For row-scaled NVFP4 this + // yields a per-column amax vector; otherwise it stays a scalar {1}. + shapes.emplace_back(amax_shape(columnwise_data_shapes[i], row_scaled_nvfp4)); } dtypes.insert(dtypes.end(), num_tensors, torch::kFloat32); alignments.insert(alignments.end(), num_tensors, 16); @@ -1244,7 +1381,7 @@ std::tuple, std::vector, bool> bulk_alloc } if (columnwise_usage) { tensor_wrapper.set_columnwise_amax(amax_columnwise_list[i].data_ptr(), DType::kFloat32, - std::vector{1}); + getTensorShape(amax_columnwise_list[i])); } tensor_cpp_list.emplace_back(std::move(tensor_wrapper)); @@ -1681,15 +1818,31 @@ void split_quantize_nvfp4_impl(const TensorWrapper &input, auto stream = at::cuda::getCurrentCUDAStream(); #endif + // The grouped Hadamard transform kernels are implemented for the SM100 family + // only. On other architectures, where + // NVFP4Quantizer::is_eligible_for_rht_cast_fusion is false as well, quantize + // each split on its own instead. That takes the generic unfused RHT path. + const int sm = transformer_engine::cuda::sm_arch(); + const bool grouped_rht_supported = sm >= 100 && sm <= 110; + // Perform multi-tensor quantization NVTE_SCOPED_GIL_RELEASE({ #ifndef USE_ROCM if (quantizer.with_rht) { // Quantize row-wise data, RHT+quantize column-wise data // Check that config is supported NVTE_CHECK(input.dtype() == DType::kBFloat16, "RHT is only supported for bfloat16 input"); - // Fuse the rowwise and colwise into one when the kernel is ready - split_quantize_nvfp4_impl_with_rht_helper(input, input_list, output_list, split_sections, - quantizers, stream); + if (grouped_rht_supported) { + // Fuse the rowwise and colwise into one when the kernel is ready + split_quantize_nvfp4_impl_with_rht_helper(input, input_list, output_list, split_sections, + quantizers, stream); + } else { + for (size_t i = 0; i < num_tensors; ++i) { + if (input_list[i].numel() == 0) { + continue; + } + quantizers[i]->quantize(input_list[i], output_list[i], std::nullopt); + } + } } else { // NVFP4 quantize // Fuse the rowwise and colwise into one when the kernel is ready split_quantize_nvfp4_impl_helper(input, input_list, output_list, split_sections, quantizers, diff --git a/transformer_engine/pytorch/csrc/extensions/ep.cpp b/transformer_engine/pytorch/csrc/extensions/ep.cpp index 118f14a01f..cef489dbbb 100644 --- a/transformer_engine/pytorch/csrc/extensions/ep.cpp +++ b/transformer_engine/pytorch/csrc/extensions/ep.cpp @@ -73,8 +73,14 @@ NVTECommWindow maybe_make_window(const at::Tensor& t) { NVTE_CHECK(nccl_sm != nullptr, "Symm-mem backend mismatch: expected NCCLSymmetricMemory. Set the backend to " "\"NCCL\" before allocating EP payload buffers."); - return NVTECommWindow{static_cast(nccl_sm->get_window()), - static_cast(nccl_sm->get_offset())}; + // rendezvous resolves ``t`` by its storage base, so get_offset() is the allocation's offset in + // the NCCL window. Add ``t``'s own storage offset so a slice/view of a symm-mem allocation + // (e.g. the scale region carved from a shared recv buffer) resolves to its true position in the + // window rather than the allocation base. + const uint64_t offset = + static_cast(nccl_sm->get_offset()) + + static_cast(t.storage_offset()) * static_cast(t.element_size()); + return NVTECommWindow{static_cast(nccl_sm->get_window()), offset}; #else (void)t; return kNoWindow; @@ -114,6 +120,34 @@ DType check_topk_idx_dtype(at::Tensor topk_idx) { using Shape = std::vector; +// EP block scaling supports only E4M3 MXFP8 today. A future block-scaled recipe (e.g. NVFP4) +// would also set is_scaled but carry a non-FP8 token dtype, so key the guard on the FP8 dtype. +bool is_mxfp8_scaled(bool is_scaled, const at::Tensor& data) { + return is_scaled && data.scalar_type() == at::kFloat8_e4m3fn; +} + +// Validate the scale-inverse pair of a block-scaled EP op and return the scale column count +// (hidden/block). Both scales must be 2D contiguous with matching cols dividing H and numels +// equal to their row counts times cols; the recv scale must be symm-mem-backed under zero-copy. +size_t check_mxfp8_scale_pair(const at::Tensor& send_scale, const at::Tensor& recv_scale, + size_t send_rows, size_t recv_rows, size_t H, const char* recv_name) { + NVTE_CHECK(send_scale.dim() >= 2 && recv_scale.dim() >= 2, + "scale-inverses must be at least 2D [., H/block]"); + NVTE_CHECK(send_scale.is_contiguous() && recv_scale.is_contiguous(), + "scale-inverses must be contiguous"); + const size_t sc_cols = static_cast(send_scale.size(-1)); + NVTE_CHECK(sc_cols > 0 && H % sc_cols == 0, "scale cols (", sc_cols, + ") must be a non-zero divisor of hidden (", H, ")"); + NVTE_CHECK(static_cast(recv_scale.size(-1)) == sc_cols, + "recv scale cols must match send scale cols"); + NVTE_CHECK(static_cast(send_scale.numel()) == send_rows * sc_cols, + "send scale numel must equal rows * cols"); + NVTE_CHECK(static_cast(recv_scale.numel()) == recv_rows * sc_cols, + "recv scale numel must equal rows * cols"); + check_symm_mem_required(recv_scale, recv_name); + return sc_cols; +} + } // namespace bool ep_get_zero_copy() { return g_zero_copy_enabled.load(std::memory_order_relaxed); } @@ -125,7 +159,7 @@ bool ep_get_zero_copy() { return g_zero_copy_enabled.load(std::memory_order_rela void ep_initialize(uintptr_t comm_ptr, const std::string& group_name, int64_t num_experts, int64_t max_tokens_per_rank, int64_t max_recv_tokens_per_rank, int64_t hidden_dim, int64_t max_num_sms, pybind11::object max_token_dtype, - bool zero_copy) { + bool zero_copy, int64_t num_topk, bool drop_on_overflow) { NVTE_CHECK(!group_name.empty(), "group_name must be non-empty (used for symm-mem lookup)"); NVTE_CHECK(comm_ptr != 0, "comm_ptr must be non-null (torch NCCL host comm pointer)"); NVTE_CHECK(!g_ep_initialized, "ep_initialize called twice without ep_finalize"); @@ -144,6 +178,8 @@ void ep_initialize(uintptr_t comm_ptr, const std::string& group_name, int64_t nu .num_comm_sms = static_cast(max_num_sms), .max_token_dtype = static_cast(GetTransformerEngineDType(torch_dtype)), .zero_copy = zero_copy ? 1 : 0, + .num_topk = static_cast(num_topk), + .drop_on_overflow = drop_on_overflow ? 1 : 0, }; // Release the GIL only around the native init. It must stay held while pybind11 casts // the ``max_token_dtype`` object above and destroys the by-value ``pybind11::object`` @@ -186,28 +222,43 @@ int64_t ep_handle_mem_size(int64_t top_k, int64_t dispatch_output_per_expert_ali // ── Per-step ops ───────────────────────────────────────────────────────────── -void ep_prepare(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor token_counts, int64_t top_k, - int64_t dispatch_output_per_expert_alignment) { +void ep_prepare(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens_per_expert, + int64_t top_k, int64_t dispatch_output_per_expert_alignment, + at::Tensor total_recv_tokens) { auto stream = at::cuda::getCurrentCUDAStream().stream(); NVTE_CHECK(topk_idx.dim() >= 2, "topk_idx must be at least 2D [..., top_k]"); auto idx_dtype = check_topk_idx_dtype(topk_idx); + // NCCL EP requires all prepare int output counters to share a dtype. + NVTE_CHECK(tokens_per_expert.scalar_type() == at::kLong, "tokens_per_expert must be int64"); + NVTE_CHECK(total_recv_tokens.scalar_type() == at::kLong, "total_recv_tokens must be int64"); const size_t T_flat = topk_idx.numel() / topk_idx.size(-1); const size_t topk_n = static_cast(topk_idx.size(-1)); auto topk_idx_te = makeTransformerEngineTensor(topk_idx.data_ptr(), Shape{T_flat, topk_n}, idx_dtype); - auto token_counts_te = makeTransformerEngineTensor( - token_counts.data_ptr(), Shape{static_cast(token_counts.numel())}, DType::kInt32); + auto tokens_per_expert_te = makeTransformerEngineTensor( + tokens_per_expert.data_ptr(), Shape{static_cast(tokens_per_expert.numel())}, + DType::kInt64); auto handle_mem_te = makeTransformerEngineTensor( handle_mem.data_ptr(), Shape{static_cast(handle_mem.numel())}, DType::kByte); + // [1] int64 scalar recv-slot total; lets the caller size dispatch outputs + // (eager) or detect overflow past recv_capacity_per_rank (graph mode). + auto total_recv_tokens_te = makeTransformerEngineTensor( + total_recv_tokens.data_ptr(), Shape{static_cast(total_recv_tokens.numel())}, + DType::kInt64); auto layer_cfg = make_layer_cfg(top_k, dispatch_output_per_expert_alignment); - nvte_ep_prepare(handle_mem_te.data(), topk_idx_te.data(), token_counts_te.data(), - /*total_recv_tokens_per_rank=*/nullptr, &layer_cfg, stream); + nvte_ep_prepare(handle_mem_te.data(), topk_idx_te.data(), tokens_per_expert_te.data(), + total_recv_tokens_te.data(), &layer_cfg, stream); } +// tokens_scale_inv / recv_scale_inv are set only for block-scaled dispatch (for +// now MXFP8): tokens/recv_tokens carry e4m3 data and the scale tensors carry the +// unswizzled e8m0 scale-inverses [T, H/block]. Both null => bf16/fp16/fp32. void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens, - at::Tensor topk_weights, at::Tensor recv_tokens, at::Tensor recv_topk_weights) { + at::Tensor topk_weights, at::Tensor recv_tokens, at::Tensor recv_topk_weights, + std::optional tokens_scale_inv, + std::optional recv_scale_inv) { auto stream = at::cuda::getCurrentCUDAStream().stream(); NVTE_CHECK(tokens.dim() >= 2, "tokens must be at least 2D [..., H]"); NVTE_CHECK(topk_idx.dim() >= 2, "topk_idx must be at least 2D [..., top_k]"); @@ -238,16 +289,39 @@ void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens, check_symm_mem_required(recv_tokens, "recv_tokens"); check_symm_mem_required(recv_topk_weights, "recv_topk_weights"); + // Block-scaled dispatch: tokens carry e4m3 data and the scale tensors carry + // unswizzled e8m0 scale-inverses [T, H/block]. Scales ride in the tensor; the + // backend keys on the tensor's scaling mode. is_mxfp8 is split from is_scaled + // so future block-scaled recipes can reuse the scale-routing plumbing while + // building their own TE tensors; only MXFP8 is supported for now. + const bool is_scaled = tokens_scale_inv.has_value(); + const bool is_mxfp8 = is_mxfp8_scaled(is_scaled, tokens); + size_t sc_cols = 0; + if (is_scaled) { + NVTE_CHECK(recv_scale_inv.has_value(), + "recv_scale_inv must be provided together with tokens_scale_inv"); + NVTE_CHECK(is_mxfp8, "EP dispatch currently supports only E4M3 MXFP8 block scaling"); + sc_cols = check_mxfp8_scale_pair(*tokens_scale_inv, *recv_scale_inv, T_flat, recv_pr, H, + "recv_scale_inv"); + } + auto tok_dtype = GetTransformerEngineDType(tokens.scalar_type()); auto handle_mem_te = makeTransformerEngineTensor( handle_mem.data_ptr(), Shape{static_cast(handle_mem.numel())}, DType::kByte); auto topk_idx_te = makeTransformerEngineTensor(topk_idx.data_ptr(), Shape{T_flat, topk_n}, idx_dtype); - auto tokens_te = makeTransformerEngineTensor(tokens.data_ptr(), Shape{T_flat, H}, tok_dtype); + auto tokens_te = + is_mxfp8 ? makeTransformerEngineTensor(tokens.data_ptr(), Shape{T_flat, H}, tok_dtype, + nullptr, nullptr, tokens_scale_inv->data_ptr(), + Shape{T_flat, sc_cols}, NVTE_MXFP8_1D_SCALING) + : makeTransformerEngineTensor(tokens.data_ptr(), Shape{T_flat, H}, tok_dtype); auto topk_w_te = makeTransformerEngineTensor(topk_weights.data_ptr(), Shape{T_flat, topk_n}, DType::kFloat32); auto recv_tokens_te = - makeTransformerEngineTensor(recv_tokens.data_ptr(), Shape{recv_pr, H}, tok_dtype); + is_mxfp8 ? makeTransformerEngineTensor(recv_tokens.data_ptr(), Shape{recv_pr, H}, tok_dtype, + nullptr, nullptr, recv_scale_inv->data_ptr(), + Shape{recv_pr, sc_cols}, NVTE_MXFP8_1D_SCALING) + : makeTransformerEngineTensor(recv_tokens.data_ptr(), Shape{recv_pr, H}, tok_dtype); auto recv_topk_w_te = makeTransformerEngineTensor(recv_topk_weights.data_ptr(), Shape{recv_pr}, DType::kFloat32); @@ -257,6 +331,17 @@ void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens, NVTECommWindow topk_w_win = maybe_make_window(topk_weights); NVTECommWindow recv_tokens_win = maybe_make_window(recv_tokens); NVTECommWindow recv_topk_w_win = maybe_make_window(recv_topk_weights); + // Block-scaled zero-copy: the scale-inverse rides on the data tensor's window. + // Send scales (tokens_scale_inv) stay staged like the send data; recv scales + // are the one-sided write target and must be symm-mem-backed under zero-copy. + if (is_scaled) { + const NVTECommWindow tsi_win = maybe_make_window(*tokens_scale_inv); + const NVTECommWindow rsi_win = maybe_make_window(*recv_scale_inv); + tokens_win.scale_window = tsi_win.window; + tokens_win.scale_offset = tsi_win.offset; + recv_tokens_win.scale_window = rsi_win.window; + recv_tokens_win.scale_offset = rsi_win.offset; + } nvte_ep_dispatch(handle_mem_te.data(), topk_idx_te.data(), tokens_te.data(), tokens_win, topk_w_te.data(), topk_w_win, recv_tokens_te.data(), recv_tokens_win, recv_topk_w_te.data(), recv_topk_w_win, stream); @@ -332,7 +417,9 @@ void ep_dispatch_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor g_recv_t g_recv_w_win, grad_tokens_te.data(), grad_topk_w_te.data(), stream); } -void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expert_out) { +void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expert_out, + std::optional grad_scale_inv, + std::optional grad_expert_out_scale_inv) { auto stream = at::cuda::getCurrentCUDAStream().stream(); NVTE_CHECK(grad.dim() >= 2, "grad must be at least 2D [..., H]"); NVTE_CHECK(grad_expert_out.dim() >= 2, "grad_expert_out must be at least 2D [..., recv_pr, H]"); @@ -351,17 +438,51 @@ void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expe // EpBuffer-owned scatter target and must be symm-mem in zero-copy mode. check_symm_mem_required(grad_expert_out, "grad_expert_out"); + // Block-scaled (MXFP8) backward: grad/grad_expert_out carry e4m3 data and the scale + // tensors carry the unswizzled e8m0 scale-inverses [., H/block]; the reverse-direction + // dispatch forwards them like the forward path. is_mxfp8 is split from is_scaled so + // future block-scaled recipes can reuse the plumbing; only MXFP8 is supported for now. + const bool is_scaled = grad_scale_inv.has_value(); + const bool is_mxfp8 = is_mxfp8_scaled(is_scaled, grad); + size_t sc_cols = 0; + if (is_scaled) { + NVTE_CHECK(grad_expert_out_scale_inv.has_value(), + "grad_expert_out_scale_inv must be provided together with grad_scale_inv"); + NVTE_CHECK(is_mxfp8, "EP combine backward currently supports only E4M3 MXFP8 block scaling"); + sc_cols = check_mxfp8_scale_pair(*grad_scale_inv, *grad_expert_out_scale_inv, T_flat, recv_pr, + H, "grad_expert_out_scale_inv"); + } + auto g_dtype = GetTransformerEngineDType(grad.scalar_type()); auto handle_mem_te = makeTransformerEngineTensor( handle_mem.data_ptr(), Shape{static_cast(handle_mem.numel())}, DType::kByte); - auto grad_te = makeTransformerEngineTensor(grad.data_ptr(), Shape{T_flat, H}, g_dtype); + auto grad_te = is_mxfp8 + ? makeTransformerEngineTensor(grad.data_ptr(), Shape{T_flat, H}, g_dtype, + nullptr, nullptr, grad_scale_inv->data_ptr(), + Shape{T_flat, sc_cols}, NVTE_MXFP8_1D_SCALING) + : makeTransformerEngineTensor(grad.data_ptr(), Shape{T_flat, H}, g_dtype); auto grad_expert_out_te = - makeTransformerEngineTensor(grad_expert_out.data_ptr(), Shape{recv_pr, H}, g_dtype); + is_mxfp8 + ? makeTransformerEngineTensor(grad_expert_out.data_ptr(), Shape{recv_pr, H}, g_dtype, + nullptr, nullptr, grad_expert_out_scale_inv->data_ptr(), + Shape{recv_pr, sc_cols}, NVTE_MXFP8_1D_SCALING) + : makeTransformerEngineTensor(grad_expert_out.data_ptr(), Shape{recv_pr, H}, g_dtype); // grad is autograd-allocated (staged); grad_expert_out resolves to a symm-mem // window in zero-copy mode, else kNoWindow for the staged path. NVTECommWindow grad_win = maybe_make_window(grad); NVTECommWindow grad_expert_out_win = maybe_make_window(grad_expert_out); + // Block-scaled zero-copy: the scale-inverse rides on the data tensor's window, + // mirroring the forward dispatch. Send scales stay staged like the send data; + // recv scales are the one-sided write target and must be symm-mem-backed. + if (is_scaled) { + const NVTECommWindow gsi_win = maybe_make_window(*grad_scale_inv); + const NVTECommWindow gesi_win = maybe_make_window(*grad_expert_out_scale_inv); + grad_win.scale_window = gsi_win.window; + grad_win.scale_offset = gsi_win.offset; + grad_expert_out_win.scale_window = gesi_win.window; + grad_expert_out_win.scale_offset = gesi_win.offset; + } nvte_ep_combine_bwd(handle_mem_te.data(), grad_te.data(), grad_win, grad_expert_out_te.data(), grad_expert_out_win, stream); } @@ -372,19 +493,28 @@ void register_ep_bindings(pybind11::module_& m) { "Initialize the EP backend; borrows torch's NCCL comm pointed to by ``comm_ptr``.", py::arg("comm_ptr"), py::arg("group_name"), py::arg("num_experts"), py::arg("max_tokens_per_rank"), py::arg("max_recv_tokens_per_rank"), py::arg("hidden_dim"), - py::arg("max_num_sms") = 0, py::arg("max_token_dtype"), py::arg("zero_copy") = false); + py::arg("max_num_sms") = 0, py::arg("max_token_dtype"), py::arg("zero_copy") = false, + py::arg("num_topk") = 0, py::arg("drop_on_overflow") = false); m.def("ep_finalize", &ep_finalize, "Tear down the EP backend. Idempotent.", py::call_guard()); m.def("ep_get_zero_copy", &ep_get_zero_copy, "Return the current EP zero-copy toggle state."); m.def("ep_handle_mem_size", &ep_handle_mem_size, "Return the handle_mem byte size for the given layer config.", py::arg("top_k"), py::arg("dispatch_output_per_expert_alignment") = 0); - m.def("ep_prepare", &ep_prepare, "EP prepare", py::call_guard()); - m.def("ep_dispatch", &ep_dispatch, "EP dispatch", py::call_guard()); + m.def("ep_prepare", &ep_prepare, "EP prepare", py::arg("handle_mem"), py::arg("topk_idx"), + py::arg("tokens_per_expert"), py::arg("top_k"), + py::arg("dispatch_output_per_expert_alignment"), py::arg("total_recv_tokens"), + py::call_guard()); + m.def("ep_dispatch", &ep_dispatch, "EP dispatch", py::arg("handle_mem"), py::arg("topk_idx"), + py::arg("tokens"), py::arg("topk_weights"), py::arg("recv_tokens"), + py::arg("recv_topk_weights"), py::arg("tokens_scale_inv") = std::nullopt, + py::arg("recv_scale_inv") = std::nullopt, py::call_guard()); m.def("ep_combine", &ep_combine, "EP combine", py::call_guard()); m.def("ep_dispatch_bwd", &ep_dispatch_bwd, "EP dispatch backward", py::call_guard()); - m.def("ep_combine_bwd", &ep_combine_bwd, "EP combine backward", + m.def("ep_combine_bwd", &ep_combine_bwd, "EP combine backward", py::arg("handle_mem"), + py::arg("grad"), py::arg("grad_expert_out"), py::arg("grad_scale_inv") = std::nullopt, + py::arg("grad_expert_out_scale_inv") = std::nullopt, py::call_guard()); } diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index e1dcb8c68e..d17b60eff3 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -148,6 +148,7 @@ void init_grouped_tensor_extension() { } void init_extension() { + pybind11::gil_scoped_acquire gil; std::call_once(extension_init_flag, []() { init_float8_extension(); init_mxfp8_extension(); @@ -239,13 +240,19 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("group_quantize", transformer_engine::pytorch::group_quantize, py::arg("tensor"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), py::arg("last_dims") = py::none(), py::arg("tensor_offsets") = py::none(), - py::arg("noop_flag") = py::none()); + py::arg("noop_flag") = py::none(), py::arg("output") = py::none()); transformer_engine::pytorch::bind_quantize_with_amax_extensions(m); m.def("group_dequantize", transformer_engine::pytorch::group_dequantize, "Dequantize group tensor", py::arg("input"), py::arg("otype")); m.def("bgrad_group_quantize", transformer_engine::pytorch::bgrad_group_quantize, py::arg("tensor"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), py::arg("last_dims") = py::none(), py::arg("tensor_offsets") = py::none()); + m.def("group_requantize_inplace", transformer_engine::pytorch::group_requantize_inplace, + "Rebuild the columnwise copy of a rowwise-prequantized MXFP8 grouped tensor and swizzle " + "its rowwise scales for GEMM, in place", + py::arg("grouped_x"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), + py::arg("otype"), py::arg("tensor_offsets") = py::none(), + py::arg("return_dequantized") = false); m.def("bgrad_quantize", transformer_engine::pytorch::bgrad_quantize, "Compute bias gradient and quantize", py::arg("input"), py::arg("quantizer")); m.def("generic_gemm", transformer_engine::pytorch::gemm, "Compute GEMM (matrix-matrix multiply)", @@ -315,6 +322,27 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Backward of SwiGLU used in GPT OSS", py::arg("grad"), py::arg("fwd_input"), py::arg("quantizer"), py::arg("limit") = 7.0f, py::arg("alpha") = 1.702f, py::arg("glu_linear_offset") = 1.0f); + /* Scaled activation */ + m.def("scaled_swiglu", transformer_engine::pytorch::scaled_swiglu, "Scaled SwiGLU activation", + py::arg("input"), py::arg("act_scales"), py::arg("quantizer"), + py::arg("glu_interleave_size") = 0); + m.def("scaled_clamped_swiglu", transformer_engine::pytorch::scaled_clamped_swiglu, + "Scaled clamped SwiGLU activation", py::arg("input"), py::arg("act_scales"), + py::arg("quantizer"), py::arg("limit") = 7.0f, py::arg("alpha") = 1.702f, + py::arg("glu_linear_offset") = 1.0f, py::arg("glu_interleave_size") = 0); + m.def("scaled_srelu", transformer_engine::pytorch::scaled_srelu, "Scaled SReLU activation", + py::arg("input"), py::arg("act_scales"), py::arg("quantizer")); + m.def("scaled_dswiglu", transformer_engine::pytorch::scaled_dswiglu, "Scaled SwiGLU backward", + py::arg("grad"), py::arg("fwd_input"), py::arg("act_scales"), py::arg("quantizer"), + py::arg("glu_interleave_size") = 0, py::arg("compute_scale_grad") = true); + m.def("scaled_clamped_dswiglu", transformer_engine::pytorch::scaled_clamped_dswiglu, + "Scaled clamped SwiGLU backward", py::arg("grad"), py::arg("fwd_input"), + py::arg("act_scales"), py::arg("quantizer"), py::arg("limit") = 7.0f, + py::arg("alpha") = 1.702f, py::arg("glu_linear_offset") = 1.0f, + py::arg("glu_interleave_size") = 0, py::arg("compute_scale_grad") = true); + m.def("scaled_dsrelu", transformer_engine::pytorch::scaled_dsrelu, "Scaled SReLU backward", + py::arg("grad"), py::arg("fwd_input"), py::arg("act_scales"), py::arg("quantizer"), + py::arg("compute_scale_grad") = true); /* DBias + DAct fusions*/ m.def("dbias_dgelu", transformer_engine::pytorch::dbias_dgelu, "DGeLU + DBias + Quantize", py::arg("grad"), py::arg("fwd_input"), py::arg("quantizer")); diff --git a/transformer_engine/pytorch/csrc/extensions/softmax.cpp b/transformer_engine/pytorch/csrc/extensions/softmax.cpp index 3bb6a5e7b3..be976d2cc4 100644 --- a/transformer_engine/pytorch/csrc/extensions/softmax.cpp +++ b/transformer_engine/pytorch/csrc/extensions/softmax.cpp @@ -52,15 +52,20 @@ at::Tensor scaled_softmax_backward(at::Tensor output_grad_, at::Tensor softmax_r (softmax_results.scalar_type() == at::ScalarType::BFloat16), "Only fp16 and bf16 are supported"); + // Allocate a fresh output buffer so the op does not alias / mutate its + // inputs (required by `torch.library.custom_op`). + auto input_grads = + torch::empty(output_grads.sizes(), output_grads.options().requires_grad(false)); + auto output_grads_cu = makeTransformerEngineTensor(output_grads); auto softmax_results_cu = makeTransformerEngineTensor(softmax_results); + auto input_grads_cu = makeTransformerEngineTensor(input_grads); - // Produce gradients in place. nvte_scaled_softmax_backward(output_grads_cu.data(), softmax_results_cu.data(), - output_grads_cu.data(), scale_factor, + input_grads_cu.data(), scale_factor, at::cuda::getCurrentCUDAStream()); - return output_grads; + return input_grads; } at::Tensor scaled_masked_softmax_forward(at::Tensor input, at::Tensor mask, float scale_factor) { @@ -115,15 +120,20 @@ at::Tensor scaled_masked_softmax_backward(at::Tensor output_grad_, at::Tensor so (softmax_results.scalar_type() == at::ScalarType::BFloat16), "Only fp16 and bf16 are supported"); + // Allocate a fresh output buffer so the op does not alias / mutate its + // inputs (required by `torch.library.custom_op`). + auto input_grads = + torch::empty(output_grads.sizes(), output_grads.options().requires_grad(false)); + auto output_grads_cu = makeTransformerEngineTensor(output_grads); auto softmax_results_cu = makeTransformerEngineTensor(softmax_results); + auto input_grads_cu = makeTransformerEngineTensor(input_grads); - // Produce gradients in place. nvte_scaled_softmax_backward(output_grads_cu.data(), softmax_results_cu.data(), - output_grads_cu.data(), scale_factor, + input_grads_cu.data(), scale_factor, at::cuda::getCurrentCUDAStream()); - return output_grads; + return input_grads; } at::Tensor scaled_upper_triang_masked_softmax_forward(at::Tensor input, float scale_factor) { @@ -167,15 +177,20 @@ at::Tensor scaled_upper_triang_masked_softmax_backward(at::Tensor output_grads_, TORCH_CHECK(output_grads.size(1) == output_grads.size(2)); + // Allocate a fresh output buffer so the op does not alias / mutate its + // inputs (required by `torch.library.custom_op`). + auto input_grads = + torch::empty(output_grads.sizes(), output_grads.options().requires_grad(false)); + auto output_grads_cu = makeTransformerEngineTensor(output_grads); auto softmax_results_cu = makeTransformerEngineTensor(softmax_results); + auto input_grads_cu = makeTransformerEngineTensor(input_grads); - // Produce gradients in place. - nvte_scaled_upper_triang_masked_softmax_backward( - output_grads_cu.data(), softmax_results_cu.data(), output_grads_cu.data(), scale_factor, - at::cuda::getCurrentCUDAStream()); + nvte_scaled_upper_triang_masked_softmax_backward(output_grads_cu.data(), + softmax_results_cu.data(), input_grads_cu.data(), + scale_factor, at::cuda::getCurrentCUDAStream()); - return output_grads; + return input_grads; } at::Tensor scaled_aligned_causal_masked_softmax_forward(at::Tensor input, float scale_factor) { @@ -223,15 +238,20 @@ at::Tensor scaled_aligned_causal_masked_softmax_backward(at::Tensor output_grad_ (softmax_results.scalar_type() == at::ScalarType::BFloat16), "Only fp16 and bf16 are supported"); + // Allocate a fresh output buffer so the op does not alias / mutate its + // inputs (required by `torch.library.custom_op`). + auto input_grads = + torch::empty(output_grads.sizes(), output_grads.options().requires_grad(false)); + auto output_grads_cu = makeTransformerEngineTensor(output_grads); auto softmax_results_cu = makeTransformerEngineTensor(softmax_results); + auto input_grads_cu = makeTransformerEngineTensor(input_grads); - // Produce gradients in place. nvte_scaled_aligned_causal_masked_softmax_backward( - output_grads_cu.data(), softmax_results_cu.data(), output_grads_cu.data(), scale_factor, + output_grads_cu.data(), softmax_results_cu.data(), input_grads_cu.data(), scale_factor, at::cuda::getCurrentCUDAStream()); - return output_grads; + return input_grads; } } // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp index e642d7a0f4..51b85dce05 100644 --- a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp +++ b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp @@ -408,6 +408,23 @@ std::optional maybe_swizzle_grouped_tensor(GroupedTensorW tensor_offsets.data_ptr, static_cast(tensor_offsets.dtype), tensor_offsets.shape); } + // Varying per-tensor dimensions. Leaving these unset declares the grouped tensor uniform, + // which selects the uniform-shape swizzle kernel. + const auto first_dims = input.get_first_dims(); + if (first_dims.data_ptr != nullptr) { + swizzle_input.set_first_dims(first_dims.data_ptr, static_cast(first_dims.dtype), + first_dims.shape); + swizzle_output.set_first_dims(first_dims.data_ptr, static_cast(first_dims.dtype), + first_dims.shape); + } + const auto last_dims = input.get_last_dims(); + if (last_dims.data_ptr != nullptr) { + swizzle_input.set_last_dims(last_dims.data_ptr, static_cast(last_dims.dtype), + last_dims.shape); + swizzle_output.set_last_dims(last_dims.data_ptr, static_cast(last_dims.dtype), + last_dims.shape); + } + // Per-tensor logical dimensions (uniform-shape grouped tensor). const size_t num_tensors = input.num_tensors(); const auto logical_shape_nvte = input.logical_shape(); @@ -417,11 +434,21 @@ std::optional maybe_swizzle_grouped_tensor(GroupedTensorW const size_t per_tensor_last_dim = logical_shape_nvte.data[logical_shape_nvte.ndim - 1]; constexpr size_t kMxfp8BlockSize = 32; - // Output is always allocated in the per-tensor padded ("swizzle-ready") layout - // so the cuDNN grouped GEMM consumer sees the correct stride between experts. - // The swizzle kernel itself handles converting from the kernel-emitted compact - // layout (per-tensor first dim is the unpadded value) to this padded layout. - auto compute_padded_grouped_scale_shape = [&](bool rowwise) { + const bool variable_shape = first_dims.data_ptr != nullptr || last_dims.data_ptr != nullptr; + + // Output is allocated in the layout the swizzle kernel writes so its consumer sees the + // correct stride between experts. + auto compute_padded_grouped_scale_shape = [&](bool rowwise) -> std::vector { + if (variable_shape) { + // Grouped variable-shape scale storage is a concatenation of per-tensor padded + // regions whose sizes live on the device. The swizzle kernel walks input and output + // with identical per-group strides, so swizzling is size-preserving and the output + // needs exactly the input's shape. The uniform-average formula below would be wrong + // here: the average of per-group sizes is generally not tile-aligned, so its + // rounded-up total matches neither the kernel's walk nor the input's capacity. + const auto &scales = rowwise ? row_scales : col_scales; + return nvte_shape_to_vector(scales.shape); + } const size_t m = rowwise ? per_tensor_first_dim : per_tensor_last_dim; const size_t k = rowwise ? per_tensor_last_dim : per_tensor_first_dim; const size_t padded_m = ceildiv(m, size_t{128}) * 128; diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 8cbea98c15..f9e38ad468 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -19,6 +19,25 @@ namespace transformer_engine::pytorch { namespace { +/*! @brief Reject unsupported columnwise-only per-tensor FP8 quantization + * + * Per-tensor FP8 uses the legacy cast-transpose kernel whenever columnwise + * output is requested. That kernel requires both rowwise and columnwise + * output buffers, so a transpose-only output cannot be produced without an + * otherwise-unused rowwise allocation and write. Keep this limitation + * explicit until the kernel supports independently selecting its outputs. + */ +void check_per_tensor_fp8_quantize_output(const TensorWrapper& output) { + const auto rowwise_data = output.get_rowwise_data(); + const auto columnwise_data = output.get_columnwise_data(); + if (rowwise_data.data_ptr == nullptr && columnwise_data.data_ptr != nullptr) { + PyErr_SetString(PyExc_NotImplementedError, + "Columnwise-only per-tensor FP8 quantization is not implemented; " + "the cast-transpose kernel requires rowwise output storage."); + throw py::error_already_set(); + } +} + /*! @brief Resolve an optional device to a concrete CUDA device * * If no device is provided, uses the current CUDA device. @@ -618,6 +637,7 @@ void Float8Quantizer::quantize(const TensorWrapper& input, TensorWrapper& out, if (input.numel() == 0) { return; } + check_per_tensor_fp8_quantize_output(out); QuantizationConfigWrapper quant_config; if (noop_flag) { quant_config.set_noop_tensor(noop_flag->data()); @@ -978,6 +998,7 @@ void Float8CurrentScalingQuantizer::quantize_impl(const TensorWrapper& input, Te out.set_scale(nullptr, DType::kFloat32, out.defaultShape); return; } + check_per_tensor_fp8_quantize_output(out); // Quantization configs QuantizationConfigWrapper quant_config; @@ -1165,14 +1186,6 @@ std::pair Float8BlockQuantizer::create_grouped const size_t logical_last_dim) const { using namespace pybind11::literals; - // The fused grouped FP8 block-scaling path uses unconstrained FP32 scales and does not - // implement power-of-2 scaling. Reject force_pow_2_scales rather than silently ignoring it; - // the unfused per-tensor split-quantize path still honors it. - NVTE_CHECK(!force_pow_2_scales, - "Fused grouped FP8 block-scaling quantize does not support force_pow_2_scales=True. " - "Set force_pow_2_scales=False, or use the unfused split-quantize path " - "(NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=0) which supports power-of-2 scales."); - const auto tensor_offsets = resolve_grouped_tensor_offsets(num_tensors, first_dims, last_dims, precomputed_tensor_offsets, logical_first_dim, logical_last_dim); @@ -1508,6 +1521,7 @@ std::vector Float8BlockQuantizer::get_scale_shape(const std::vectordtype = quantizer.attr("dtype").cast(); + this->with_2d_quantization = quantizer.attr("with_2d_quantization").cast(); } void MXFP8Quantizer::set_quantization_params(TensorWrapper* tensor) const {} @@ -1854,6 +1868,9 @@ void MXFP8Quantizer::quantize(const TensorWrapper& input, TensorWrapper& out, if (noop_flag) { quant_config.set_noop_tensor(noop_flag->data()); } + if (this->with_2d_quantization) { + quant_config.set_mxfp8_2d_quantization(true); + } NVTE_SCOPED_GIL_RELEASE({ nvte_quantize_v2(input.data(), out.data(), quant_config, at::cuda::getCurrentCUDAStream()); }); @@ -2015,8 +2032,6 @@ std::pair NVFP4Quantizer::create_tensor( const int nvfp4_e4m3_max = this->nvfp4_e4m3_max; if (row_scaled_nvfp4) { NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 quantization requires rowwise usage."); - NVTE_CHECK(!columnwise_usage, - "Row-scaled NVFP4 quantization does not support columnwise usage."); } const auto rowwise_scale_inv_shape = get_scale_shape(shape, false); const auto columnwise_scale_inv_shape = get_scale_shape(shape, true); @@ -2051,7 +2066,8 @@ std::pair NVFP4Quantizer::create_tensor( columnwise_scale_inv_tensor = at::empty(scale_inv_shape_int64, bit8_tensor_opts); // hadamard amax kernel will zero out pointer with ZeroAmaxKernel // nvte_compute_amax_with_config will zero out the pointer if needed - amax_columnwise = at::empty({1}, bit32_tensor_opts); + const int64_t amax_cols = row_scaled_nvfp4 ? static_cast(flat_last_dim) : 1; + amax_columnwise = at::empty({amax_cols}, bit32_tensor_opts); } // Convert tensors to Python @@ -2143,7 +2159,7 @@ std::pair NVFP4Quantizer::create_tensor( out_cpp.set_columnwise_scale_inv(columnwise_scale_inv_tensor.data_ptr(), DType::kFloat8E4M3, columnwise_scale_inv_shape); out_cpp.set_columnwise_amax(amax_columnwise.data_ptr(), DType::kFloat32, - std::vector{1}); + getTensorShape(amax_columnwise)); } out_cpp.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); out_cpp.set_row_scaled_nvfp4(row_scaled_nvfp4); @@ -2342,8 +2358,6 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( const int nvfp4_e4m3_max = this->nvfp4_e4m3_max; if (row_scaled_nvfp4) { NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 quantization requires rowwise usage."); - NVTE_CHECK(!columnwise_usage, - "Row-scaled NVFP4 quantization does not support columnwise usage."); } tensor.attr("_row_scaled_nvfp4") = row_scaled_nvfp4; tensor.attr("_with_gemm_swizzled_scales") = with_gemm_swizzled_scales; @@ -2409,11 +2423,12 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( columnwise_scale_inv = at::empty(scale_inv_shape_int64, opts); tensor.attr("_columnwise_scale_inv") = *columnwise_scale_inv; } - if (!amax_columnwise) { + const int64_t amax_cols = row_scaled_nvfp4 ? static_cast(flat_last_dim) : 1; + if (!amax_columnwise || amax_columnwise->numel() != amax_cols) { const auto opts = at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); // hadamard amax kernel will zero out pointer with ZeroAmaxKernel // nvte_compute_amax_with_config will zero out the pointer if needed - amax_columnwise = at::empty({1}, opts); + amax_columnwise = at::empty({amax_cols}, opts); tensor.attr("_amax_columnwise") = *amax_columnwise; } } else { // columnwise_usage == false @@ -2449,7 +2464,7 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( out_cpp.set_columnwise_scale_inv(columnwise_scale_inv->data_ptr(), DType::kFloat8E4M3, getTensorShape(*columnwise_scale_inv)); out_cpp.set_columnwise_amax(amax_columnwise->data_ptr(), DType::kFloat32, - std::vector{1}); + getTensorShape(*amax_columnwise)); } out_cpp.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); out_cpp.set_row_scaled_nvfp4(row_scaled_nvfp4); diff --git a/transformer_engine/pytorch/csrc/type_converters.cpp b/transformer_engine/pytorch/csrc/type_converters.cpp index 84ac7bcd31..4f2beb7f55 100644 --- a/transformer_engine/pytorch/csrc/type_converters.cpp +++ b/transformer_engine/pytorch/csrc/type_converters.cpp @@ -261,12 +261,23 @@ GroupedTensorWrapper GroupedTensorFromPyTorchGroupedTensor(py::handle tensor) { } auto ret = GroupedTensorWrapper(num_tensors, logical_shape, scaling_mode); + auto get_initialized_storage_shape = [](const at::Tensor &data) { + auto shape = getTensorShape(data); + if (data.numel() == 0) { + // PyTorch may use a null pointer for a valid zero-sized allocation. TE Common reserves + // {nullptr, {0}} for uninitialized storage, so use an orientation-neutral 2D empty shape + // to distinguish explicitly provided rowwise or columnwise storage from missing storage. + shape = {0, 0}; + } + return shape; + }; + // Rowwise data if (!tensor.attr("rowwise_data").is_none()) { const auto &data = tensor.attr("rowwise_data").cast(); DType data_dtype = quantizer.is_none() ? GetTransformerEngineDType(data.scalar_type()) : quantizer_dtype; - ret.set_rowwise_data(data.data_ptr(), data_dtype, getTensorShape(data)); + ret.set_rowwise_data(data.data_ptr(), data_dtype, get_initialized_storage_shape(data)); } else if (quantizer_dtype != DType::kNumTypes) { ret.set_rowwise_data(nullptr, quantizer_dtype, std::vector{0}); } @@ -276,7 +287,7 @@ GroupedTensorWrapper GroupedTensorFromPyTorchGroupedTensor(py::handle tensor) { const auto &data = tensor.attr("columnwise_data").cast(); DType data_dtype = quantizer.is_none() ? GetTransformerEngineDType(data.scalar_type()) : quantizer_dtype; - ret.set_columnwise_data(data.data_ptr(), data_dtype, getTensorShape(data)); + ret.set_columnwise_data(data.data_ptr(), data_dtype, get_initialized_storage_shape(data)); } else if (quantizer_dtype != DType::kNumTypes) { ret.set_columnwise_data(nullptr, quantizer_dtype, std::vector{0}); } diff --git a/transformer_engine/pytorch/custom_recipes/gemm.py b/transformer_engine/pytorch/custom_recipes/gemm.py index 3d1e1cc43e..c2901c8023 100644 --- a/transformer_engine/pytorch/custom_recipes/gemm.py +++ b/transformer_engine/pytorch/custom_recipes/gemm.py @@ -4,18 +4,34 @@ """GEMM API that enables custom GEMM logic for custom quantization recipes.""" +import dataclasses +import enum from typing import Iterable, Optional import torch -from transformer_engine.pytorch.custom_recipes.quantization import ( - MMParams, - GEMMType, -) from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage, Quantizer from transformer_engine.pytorch.tensor.utils import is_custom +@enum.unique +class GEMMType(enum.Enum): + """Type of GEMM operation being performed.""" + + FPROP = "fprop" + DGRAD = "dgrad" + WGRAD = "wgrad" + + +@dataclasses.dataclass(frozen=True) +class MMParams: + """Matrix multiplication parameters.""" + + out_dtype: torch.dtype | None = None + # Use split accumulator for more accurate FP8 GEMM + use_split_accumulator: bool = True + + def custom_gemm( A: QuantizedTensorStorage, B: QuantizedTensorStorage, diff --git a/transformer_engine/pytorch/custom_recipes/quantization.py b/transformer_engine/pytorch/custom_recipes/quantization.py deleted file mode 100644 index 85920f5032..0000000000 --- a/transformer_engine/pytorch/custom_recipes/quantization.py +++ /dev/null @@ -1,29 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Quantization API for experimental middleware between Transformer Engine and Kitchen.""" - -from __future__ import annotations -import dataclasses -import enum - -import torch - - -@enum.unique -class GEMMType(enum.Enum): - """Type of GEMM operation being performed.""" - - FPROP = "fprop" - DGRAD = "dgrad" - WGRAD = "wgrad" - - -@dataclasses.dataclass(frozen=True) -class MMParams: - """Matrix multiplication parameters.""" - - out_dtype: torch.dtype | None = None - # Use split accumulator for more accurate FP8 GEMM - use_split_accumulator: bool = True diff --git a/transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py b/transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py deleted file mode 100644 index e88adbe4cc..0000000000 --- a/transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py +++ /dev/null @@ -1,270 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -""" -Quantizer factory examples. - -Demonstrates how to use the ``CustomRecipe`` + ``qfactory`` interface to apply -*different* quantization recipes to different module/tensor types/instances within the same model. - -Usage:: - - from transformer_engine.common.recipe import CustomRecipe - from transformer_engine.pytorch.quantization import autocast - from transformer_engine.pytorch.custom_recipes.quantization_factory_examples import ( - nvfp4_linear_mxfp8_grouped_linear_factory, - nvfp4_linear_fp8_dpa_factory, - nvfp4_linear_mxfp8_dpa_factory, - ) - - # Mixed module types: NVFP4 for Linear, MXFP8 for GroupedLinear - recipe = CustomRecipe(qfactory=nvfp4_linear_mxfp8_grouped_linear_factory) - with autocast(recipe=recipe): - output = model(input) - - # NVFP4 for Linear, FP8 current-scaling + delayed-scaling for DPA - recipe = CustomRecipe(qfactory=nvfp4_linear_fp8_dpa_factory, fp8_dpa=True) - with autocast(recipe=recipe): - output = model(input) - - # NVFP4 for Linear, MXFP8 for DPA - recipe = CustomRecipe(qfactory=nvfp4_linear_mxfp8_dpa_factory, fp8_dpa=True) - with autocast(recipe=recipe): - output = model(input) -""" - -from __future__ import annotations - -from typing import Optional - -from transformer_engine.pytorch.quantization import QuantizerRole -from ..constants import DType - - -def nvfp4_linear_mxfp8_grouped_linear_factory( - role: Optional[QuantizerRole], -): - """Quantizer factory: NVFP4 for ``Linear``, MXFP8 for ``GroupedLinear``. - - Dispatch logic: - * ``role.module_type == "grouped_linear"`` -> MXFP8 (E4M3, block-32) - * everything else (``"linear"`` or unknown) -> NVFP4 (E2M1) - - NVFP4 settings follow the built-in ``NVFP4BlockScaling`` defaults: - * Weights: 2D quantization (16x16), no RHT, no stochastic rounding - * Inputs: 1D quantization, RHT enabled, no stochastic rounding - * Grads: 1D quantization, RHT enabled, stochastic rounding enabled - """ - is_grouped_linear = role is not None and role.module_type == "grouped_linear" - - if is_grouped_linear: - return _make_mxfp8_quantizer() - - return _make_nvfp4_quantizer(role) - - -def _make_mxfp8_quantizer(): - """Return an MXFP8 quantizer with default settings (E4M3, block-32, E8M0 scales).""" - from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer - - return MXFP8Quantizer( - fp8_dtype=DType.kFloat8E4M3, - ) - - -def _make_nvfp4_quantizer(role: Optional[QuantizerRole]): - """Return an NVFP4 quantizer configured per tensor role. - - Mirrors :class:`NVFP4BlockScaling` recipe defaults. - """ - from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer - - is_linear = role is not None and role.module_type == "linear" - is_weight = is_linear and role.tensor_type == "weight" - is_grad = is_linear and role.tensor_type == "grad_output" - - if is_weight: - return NVFP4Quantizer( - fp4_dtype=DType.kFloat4E2M1, - with_rht=False, - with_post_rht_amax=False, - with_2d_quantization=True, - stochastic_rounding=False, - with_random_sign_mask=True, - ) - - if is_grad: - return NVFP4Quantizer( - fp4_dtype=DType.kFloat4E2M1, - rowwise=True, - columnwise=True, - with_rht=True, - with_post_rht_amax=True, - with_2d_quantization=False, - stochastic_rounding=True, - with_random_sign_mask=True, - ) - - return NVFP4Quantizer( - fp4_dtype=DType.kFloat4E2M1, - rowwise=True, - columnwise=True, - with_rht=True, - with_post_rht_amax=True, - with_2d_quantization=False, - stochastic_rounding=False, - with_random_sign_mask=True, - ) - - -def nvfp4_linear_fp8_dpa_factory( - role: Optional[QuantizerRole], -): - """Quantizer factory: NVFP4 for ``Linear``, mixed FP8 for ``DotProductAttention``. - - This factory demonstrates how to use ``CustomRecipe`` with ``fp8_dpa=True`` - to combine NVFP4 quantization for linear layers with FP8 attention. - - DPA tensor types (``role.module_type == "dpa"``): - - =========== ============================================================ - tensor_type Description - =========== ============================================================ - ``"qkv"`` Query, Key, Value inputs to the first attention GEMM - ``"s"`` Softmax output (S = softmax(Q·K^T)), fed into the second GEMM - ``"o"`` Attention output (O = S·V) - ``"do"`` Gradient of the attention output (dO), backward input - ``"dp"`` Gradient of the softmax output (dP = dO·V^T), backward - ``"dqkv"`` Gradient flowing back to Q, K, V - =========== ============================================================ - - Dispatch logic: - * ``role.module_type == "dpa"`` with ``tensor_type in ("s", "dp")`` - -> FP8 delayed scaling (stateful amax tracking) - * ``role.module_type == "dpa"`` (QKV, dO) - -> FP8 current scaling (E4M3) - * DPA boundary hints (``"dpa_output"`` / ``"dpa_grad_input"`` in ``role.name``) - -> FP8 current scaling placeholder. The fused attention kernel requires - FP8-compatible quantizers in all DPA slots, even when the output is - produced in BF16 (``fp8_mha=False``). DPA emits these hint-only roles - (with empty ``module_type`` and ``tensor_type``) when the downstream - consumer is unknown. - * everything else (``"linear"`` / ``"grouped_linear"`` / ``None``) - -> NVFP4 (E2M1), configured per tensor role - - Usage:: - - from transformer_engine.common.recipe import CustomRecipe - from transformer_engine.pytorch.quantization import autocast - from transformer_engine.pytorch.custom_recipes.quantization_factory_examples import ( - nvfp4_linear_fp8_dpa_factory, - ) - - recipe = CustomRecipe( - qfactory=nvfp4_linear_fp8_dpa_factory, - fp8_dpa=True, - ) - with autocast(recipe=recipe): - output = model(input) - """ - from transformer_engine.pytorch.quantization import DelayedScalingRequest - from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer - - is_dpa = role is not None and role.module_type == "dpa" - is_softmax_or_dp = is_dpa and role.tensor_type in ("s", "dp") - - if is_softmax_or_dp: - return DelayedScalingRequest() - - if is_dpa: - return Float8CurrentScalingQuantizer( - fp8_dtype=DType.kFloat8E4M3, - device="cuda", - ) - - # DPA boundary slots (O output / dQKV grad-input): the fused attention - # kernel only supports FP8 quantizers here, regardless of the linear recipe. - is_dpa_boundary = ( - role is not None - and not role.module_type - and ("dpa_output" in role.name or "dpa_grad_input" in role.name) - ) - if is_dpa_boundary: - return Float8CurrentScalingQuantizer( - fp8_dtype=DType.kFloat8E4M3, - device="cuda", - ) - - return _make_nvfp4_quantizer(role) - - -def nvfp4_linear_mxfp8_dpa_factory( - role: Optional[QuantizerRole], -): - """Quantizer factory: NVFP4 for ``Linear``, MXFP8 for ``DotProductAttention``. - - Mirrors the documented "NVFP4 linear + MXFP8 attention" combo from - :mod:`transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention` - (see the recipe-combination table at the top of that module). With - ``CustomRecipe`` the per-tensor decision is made directly here, so the - ``NVTE_DPA_FP8_RECIPE="MXFP8BlockScaling"`` env override that the - built-in recipes would otherwise need is unnecessary. - - DPA tensor types (``role.module_type == "dpa"``): - - =========== ============================================================ - tensor_type Description - =========== ============================================================ - ``"qkv"`` Query, Key, Value inputs to the first attention GEMM - ``"s"`` Softmax output (S = softmax(Q·K^T)), fed into the second GEMM - ``"o"`` Attention output (O = S·V) - ``"do"`` Gradient of the attention output (dO), backward input - ``"dp"`` Gradient of the softmax output (dP = dO·V^T), backward - ``"dqkv"`` Gradient flowing back to Q, K, V - =========== ============================================================ - - Dispatch logic: - * ``role.module_type == "dpa"`` -> MXFP8 (E4M3, block-32) - The MXFP8 fused-attention kernel handles the S/dP slots - internally, so any quantizer returned for those roles is later - nulled out by ``get_attention_quantizers``. Returning MXFP8 is - the simplest valid choice. - * DPA boundary hints (``"dpa_output"`` / ``"dpa_grad_input"`` in - ``role.name``) -> MXFP8 placeholder. The fused attention kernel - requires FP8-compatible quantizers in all DPA slots. - * everything else (``"linear"`` / ``"grouped_linear"`` / ``None``) - -> NVFP4 (E2M1), configured per tensor role. - - Usage:: - - from transformer_engine.common.recipe import CustomRecipe - from transformer_engine.pytorch.quantization import autocast - from transformer_engine.pytorch.custom_recipes.quantization_factory_examples import ( - nvfp4_linear_mxfp8_dpa_factory, - ) - - recipe = CustomRecipe( - qfactory=nvfp4_linear_mxfp8_dpa_factory, - fp8_dpa=True, - ) - with autocast(recipe=recipe): - output = model(input) - """ - is_dpa = role is not None and role.module_type == "dpa" - if is_dpa: - return _make_mxfp8_quantizer() - - # DPA boundary slots (O output / dQKV grad-input): emitted by DPA with - # empty `module_type` and a `name` like ".dpa_output". The fused - # attention kernel requires an FP8-compatible quantizer here even when - # the downstream consumer is unknown. - is_dpa_boundary = ( - role is not None - and not role.module_type - and ("dpa_output" in role.name or "dpa_grad_input" in role.name) - ) - if is_dpa_boundary: - return _make_mxfp8_quantizer() - - return _make_nvfp4_quantizer(role) diff --git a/transformer_engine/pytorch/custom_recipes/quantization_recipes_base.py b/transformer_engine/pytorch/custom_recipes/quantizer_factories.py similarity index 78% rename from transformer_engine/pytorch/custom_recipes/quantization_recipes_base.py rename to transformer_engine/pytorch/custom_recipes/quantizer_factories.py index 45febaf413..250ce8c939 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_recipes_base.py +++ b/transformer_engine/pytorch/custom_recipes/quantizer_factories.py @@ -3,22 +3,22 @@ # See LICENSE for license information. """ -Quantizer factory examples using real silicon quantizers. +Quantizer factories that mirror Transformer Engine's built-in recipes. -Each factory below replicates the behaviour of built-in TE recipe but via the -``CustomRecipe`` + ``qfactory`` interface. This is useful when you want to -start from a known-good recipe and then selectively override quantizer settings -for specific layers / tensor types. +For TE ``Linear`` and ``GroupedLinear`` roles, each factory below mirrors the +nominal defaults of a built-in recipe through the ``CustomRecipe`` + +``qfactory`` interface. This provides a built-in-equivalent starting point for +selectively overriding quantizer settings for specific layers or tensor types. Usage (any factory):: from transformer_engine.common.recipe import CustomRecipe from transformer_engine.pytorch.quantization import autocast - from transformer_engine.pytorch.custom_recipes.quantization_recipes_base import ( - nvfp4_quantizer_factory, + from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + nvfp4_factory, ) - recipe = CustomRecipe(qfactory=nvfp4_quantizer_factory) + recipe = CustomRecipe(qfactory=nvfp4_factory) with autocast(recipe=recipe): output = model(input) """ @@ -32,7 +32,25 @@ from ..constants import DType -def delayed_scaling_quantizer_factory( +def high_precision_factory( + role: Optional[QuantizerRole], # pylint: disable=unused-argument +) -> "IdentityQuantizer": + """Factory that runs all GEMMs in high precision (no quantization). + + Returns an :class:`IdentityQuantizer` for every slot, so no tensor is + quantized. This is the simplest base factory and a good starting point to + branch from: keep most roles in high precision and selectively override the + ones you want to quantize. + + Dispatch logic: + * every role -> ``IdentityQuantizer`` (no quantization) + """ + from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer + + return IdentityQuantizer() + + +def delayed_scaling_factory( role: Optional[QuantizerRole], # pylint: disable=unused-argument ) -> "DelayedScalingRequest": """Factory that mirrors :class:`DelayedScaling` recipe defaults. @@ -51,7 +69,7 @@ def delayed_scaling_quantizer_factory( return DelayedScalingRequest(fp8_format=Format.HYBRID) -def current_scaling_quantizer_factory( +def current_scaling_factory( role: Optional[QuantizerRole], ) -> "Float8CurrentScalingQuantizer": """Factory that mirrors :class:`Float8CurrentScaling` recipe defaults. @@ -69,27 +87,12 @@ def current_scaling_quantizer_factory( return Float8CurrentScalingQuantizer( fp8_dtype=fp8_dtype, device=torch.device("cuda"), - force_pow_2_scales=False, # constrain scale to powers of 2 - amax_epsilon=0.0, # clamp amax from below to avoid div-by-zero + force_pow_2_scales=False, + amax_epsilon=0.0, ) -def mxfp8_quantizer_factory( - role: Optional[QuantizerRole], # pylint: disable=unused-argument -) -> "MXFP8Quantizer": - """Factory that mirrors :class:`MXFP8BlockScaling` recipe defaults. - - * E4M3 by default for all tensors - * Block size 32, power-of-2 (E8M0) scales - """ - from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer - - return MXFP8Quantizer( - fp8_dtype=DType.kFloat8E4M3, - ) - - -def float8_block_scaling_quantizer_factory( +def float8_block_scaling_factory( role: Optional[QuantizerRole], ) -> "Float8BlockQuantizer": """Factory that mirrors :class:`Float8BlockScaling` recipe defaults. @@ -113,13 +116,28 @@ def float8_block_scaling_quantizer_factory( fp8_dtype=DType.kFloat8E4M3, rowwise=True, columnwise=True, - amax_epsilon=0.0, # clamp amax from below to avoid div-by-zero + amax_epsilon=0.0, force_pow_2_scales=True, block_scaling_dim=block_scaling_dim, # 1 = 1D (1×128), 2 = 2D (128×128) ) -def nvfp4_quantizer_factory( +def mxfp8_factory( + role: Optional[QuantizerRole], # pylint: disable=unused-argument +) -> "MXFP8Quantizer": + """Factory that mirrors :class:`MXFP8BlockScaling` recipe defaults. + + * E4M3 by default for all tensors + * Block size 32, power-of-2 (E8M0) scales + """ + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + return MXFP8Quantizer( + fp8_dtype=DType.kFloat8E4M3, + ) + + +def nvfp4_factory( role: Optional[QuantizerRole], ) -> "NVFP4Quantizer": """Factory that mirrors :class:`NVFP4BlockScaling` recipe defaults. @@ -156,8 +174,6 @@ def nvfp4_quantizer_factory( if is_grad: return NVFP4Quantizer( fp4_dtype=DType.kFloat4E2M1, - rowwise=True, - columnwise=True, with_rht=True, with_post_rht_amax=True, with_2d_quantization=False, @@ -168,8 +184,6 @@ def nvfp4_quantizer_factory( # For input and unknown roles return NVFP4Quantizer( fp4_dtype=DType.kFloat4E2M1, - rowwise=True, - columnwise=True, with_rht=True, with_post_rht_amax=True, with_2d_quantization=False, diff --git a/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py new file mode 100644 index 0000000000..0d5fc0c441 --- /dev/null +++ b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py @@ -0,0 +1,431 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +""" +Example quantizer factories for custom and mixed quantization recipes. + +A collection of composed/mixed-recipe factories. They demonstrate how to use +the ``CustomRecipe`` + ``qfactory`` interface to apply *different* quantization +recipes to different module/tensor types/instances within the same model. +Factories may return native Transformer Engine quantizers, custom quantizers, +or ``HybridQuantizer`` instances when tensor directions should use different +representations or sources. + +Within the Linear/GroupedLinear and RL-oriented families, examples are roughly +ordered by increasingly aggressive forward quantization. This is an +organizational convention, not an expected accuracy or performance ranking. +These factories demonstrate what can be composed; they are not necessarily +tuned for end-to-end performance. + +When forward operands (inputs and weights) are quantized and use a different +representation in backward, consider setting +``columnwise_source="rowwise_dequantized"`` on their hybrid quantizers. This +applies whether backward uses another low-precision format or high precision. +It constructs backward operands from the value obtained during forward +quantization, improving forward/backward representation consistency. This +source choice should not be applied to gradient tensors. +Note: dequantization does not recover information discarded during forward quantization. + +Organization: + * Pre-training-oriented recipes: Favor more precision on the forward pass. + * RL-oriented recipes: Favor more precision in backward GEMMs. + * Linear + attention recipes: factories that also cover ``DotProductAttention`` + roles and require ``CustomRecipe(..., fp8_dpa=True)``. + +.. warning:: + + Use these with caution. These are **not** official, supported recipes + provided by Transformer Engine -- they are illustrative examples meant to + inspire your own experiments, not drop-in production defaults. Most include + a motivating rationale in their per-factory docstrings, but they have not + been broadly validated for accuracy, convergence, or performance across + models and hardware. Treat them as starting points: benchmark and verify on + your own workload before relying on any of them. + +Usage:: + + from transformer_engine.common.recipe import CustomRecipe + from transformer_engine.pytorch.quantization import autocast + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + mxfp8_fwd_nvfp4_bwd_factory, + nvfp4_linear_fp8_dpa_factory, + ) + + # Linear-only recipe (no attention quantization): the qfactory is the only knob. + recipe = CustomRecipe(qfactory=mxfp8_fwd_nvfp4_bwd_factory) + with autocast(recipe=recipe): + output = model(input) + + # Recipe that also quantizes DotProductAttention: set ``fp8_dpa=True`` so the + # attention GEMMs request quantizers from the factory (DPA roles) too. + recipe = CustomRecipe(qfactory=nvfp4_linear_fp8_dpa_factory, fp8_dpa=True) + with autocast(recipe=recipe): + output = model(input) + + # The other factories in this module follow the same two patterns; see their + # docstrings for the exact per-role dispatch. +""" + +from __future__ import annotations + +from typing import Optional + +from transformer_engine.pytorch.quantization import QuantizerRole +from ..constants import DType +from .quantizer_factories import mxfp8_factory, nvfp4_factory + +# ----------------------------------------------------------------------------- +# Pre-training-Oriented Recipes +# ----------------------------------------------------------------------------- + + +def high_precision_fwd_mxfp8_bwd_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: high-precision forward, MXFP8 backward. + + Dispatch logic: + * ``grad_output`` -> MXFP8 (E4M3, block-32) + * everything else -> ``Hybrid(rowwise=IdentityQuantizer, columnwise=MXFP8)`` + """ + from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer + from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer + + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type == "grad_output": + return mxfp8_factory(role) + + # fprop consumes rowwise high precision; dgrad / wgrad consume columnwise MXFP8. + return HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=mxfp8_factory(role), + ) + + +def _plain_nvfp4_quantizer(*, row_scaled_nvfp4: bool = False): + """NVFP4 quantizer without RHT, stochastic rounding, or 2D scaling.""" + from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer + + return NVFP4Quantizer( + fp4_dtype=DType.kFloat4E2M1, + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=False, + stochastic_rounding=False, + row_scaled_nvfp4=row_scaled_nvfp4, + ) + + +def mxfp8_fwd_nvfp4_bwd_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: MXFP8 forward, NVFP4 backward. + + Per-GEMM format consumption: + * fprop: ``weight.row(MXFP8) x input.row(MXFP8)`` + * dgrad: ``weight.col(NVFP4) x grad_output.row(NVFP4)`` + * wgrad: ``input.col(NVFP4) x grad_output.col(NVFP4)`` + + Every backward operand uses 1D NVFP4 scaling. Inputs and gradients mirror + :func:`nvfp4_factory` semantics: RHT is applied only to the columnwise + representations consumed by wgrad, and gradients use stochastic rounding. + The dgrad weight uses plain 1D NVFP4 without RHT or stochastic rounding. + + The backward weight representation consumed by dgrad is quantized to + NVFP4 from the dequantized MXFP8 forward weight. In ``HybridQuantizer`` + terms, the weight uses ``columnwise_source="rowwise_dequantized"``. The + backward input representation consumed by wgrad remains quantized directly + from the original high-precision input with ``columnwise_source="original"``. + """ + from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer + + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type == "input": + return HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=nvfp4_factory(role), + columnwise_source="original", + ) + if is_linear and role.tensor_type == "weight": + return HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=_plain_nvfp4_quantizer(), + columnwise_source="rowwise_dequantized", + ) + if is_linear and role.tensor_type == "grad_output": + return nvfp4_factory(role) + return mxfp8_factory(role) + + +def nvfp4_1d_weight_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: NVFP4 recipe with 1D weight scaling. + + Dispatch logic: + * ``linear`` / ``grouped_linear`` ``weight`` -> + ``Hybrid(rowwise=plain 1D NVFP4, columnwise=plain 1D NVFP4, + columnwise_source="rowwise_dequantized")`` + * everything else -> :func:`nvfp4_factory` + + The backward weight representation (``W.T``) is quantized to NVFP4 from + the dequantized NVFP4 forward weight rather than directly from the original + high-precision weight. In ``HybridQuantizer`` terms, this source choice is + expressed with ``columnwise_source="rowwise_dequantized"``. + + All non-weight roles keep the standard NVFP4 factory behavior, including RHT + for inputs and stochastic rounding for gradients. The weight override uses + plain 1D NVFP4 in both directions: no RHT, stochastic rounding, row-scaled + activations, or 2D weight scaling. + """ + from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer + + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type == "weight": + return HybridQuantizer( + rowwise_quantizer=_plain_nvfp4_quantizer(), + columnwise_quantizer=_plain_nvfp4_quantizer(), + columnwise_source="rowwise_dequantized", + ) + return nvfp4_factory(role) + + +# ----------------------------------------------------------------------------- +# RL-Oriented Recipes +# ----------------------------------------------------------------------------- + + +def mxfp8_fwd_high_precision_bwd_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: MXFP8 forward, high-precision backward. + + This expresses the linear/grouped-linear equivalent of + ``backward_override="dequantized"`` through per-direction quantizers: + + * ``input`` / ``weight`` -> + ``Hybrid(rowwise=MXFP8, columnwise=Identity, columnwise_source="rowwise_dequantized")`` + * ``grad_output`` -> ``IdentityQuantizer`` + * everything else -> MXFP8 + + The backward input and weight representations are high-precision values + dequantized from the MXFP8 forward representations rather than the original + high-precision tensors. In ``HybridQuantizer`` terms, this source choice is + expressed with ``columnwise_source="rowwise_dequantized"``. The gradient + output independently remains in high precision. + + This recipe targets RL-style training use cases and is motivated by + NVIDIA/TransformerEngine#2644, where ``backward_override="dequantized"`` + was introduced: + https://github.com/NVIDIA/TransformerEngine/pull/2644 + """ + from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer + from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer + + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight"): + return HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + if is_linear and role.tensor_type == "grad_output": + return IdentityQuantizer() + return mxfp8_factory(role) + + +def nvfp4_row_scaled_fwd_mxfp8_bwd_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: row-scaled NVFP4 forward, MXFP8 backward. + + This RL-related recipe is inspired by the Composer 2 MoE grouped-GEMM + recipe described in arXiv:2603.24477. + + Derived from the report: Composer 2 describes row-scaled NVFP4 for the MoE + forward pass and standard MXFP8 for the MoE backward pass. This factory maps + that format split onto ``GroupedLinear`` roles. + + Assumed here: regular non-MoE ``Linear`` layers use the MXFP8 fallback. The + public report does not specify the precision used for non-MoE linears. + + Dispatch logic: + + * ``GroupedLinear`` ``input`` -> + ``Hybrid(rowwise=row-scaled NVFP4, columnwise=MXFP8, + columnwise_source="rowwise_dequantized")`` + * ``GroupedLinear`` ``weight`` -> + ``Hybrid(rowwise=plain NVFP4, columnwise=MXFP8, + columnwise_source="rowwise_dequantized")`` + * regular ``Linear`` -> MXFP8 + * ``grad_output`` -> MXFP8 + * everything else -> MXFP8 + + Row-scaled NVFP4 is fprop-only, so the forward NVFP4 quantizers avoid RHT, + stochastic rounding, and 2D scaling. The backward input and weight + representations are quantized to MXFP8 from the dequantized NVFP4 forward + representations rather than directly from the original high-precision + tensors. In ``HybridQuantizer`` terms, this source choice is expressed with + ``columnwise_source="rowwise_dequantized"``. To use the original tensors + instead, use ``columnwise_source="original"``. + + Composer 2 Technical Report: + https://arxiv.org/abs/2603.24477 + """ + from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer + + is_grouped_linear = role is not None and role.module_type == "grouped_linear" + is_linear = role is not None and role.module_type == "linear" + if is_grouped_linear and role.tensor_type == "input": + return HybridQuantizer( + rowwise_quantizer=_plain_nvfp4_quantizer(row_scaled_nvfp4=True), + columnwise_quantizer=mxfp8_factory(role), + columnwise_source="rowwise_dequantized", + ) + if is_grouped_linear and role.tensor_type == "weight": + return HybridQuantizer( + rowwise_quantizer=_plain_nvfp4_quantizer(), + columnwise_quantizer=mxfp8_factory(role), + columnwise_source="rowwise_dequantized", + ) + if is_grouped_linear and role.tensor_type == "grad_output": + return mxfp8_factory(role) + if is_linear: + return mxfp8_factory(role) + return mxfp8_factory(role) + + +def nvfp4_row_scaled_fwd_high_precision_bwd_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: row-scaled NVFP4 forward, high-precision backward. + + This expresses a linear/grouped-linear variant of + ``NVFP4BlockScaling(row_scaled_activation=True, + backward_override="dequantized")`` through per-direction quantizers: + + * ``input`` -> + ``Hybrid(rowwise=row-scaled NVFP4, columnwise=Identity, + columnwise_source="rowwise_dequantized")`` + * ``weight`` -> + ``Hybrid(rowwise=plain NVFP4, columnwise=Identity, + columnwise_source="rowwise_dequantized")`` + * ``grad_output`` -> ``IdentityQuantizer`` + * everything else -> plain NVFP4 + + Row-scaled NVFP4 is fprop-only, so the forward quantizers avoid RHT, + stochastic rounding, and 2D scaling. + + The backward input and weight representations are high-precision values + dequantized from the NVFP4 forward representations rather than the original + high-precision tensors. In ``HybridQuantizer`` terms, this source choice is + expressed with ``columnwise_source="rowwise_dequantized"``. The gradient + output independently remains in high precision. + + This recipe targets RL-style training use cases and builds on + NVIDIA/TransformerEngine#2931, which introduced row-scaled NVFP4: + https://github.com/NVIDIA/TransformerEngine/pull/2931 + """ + from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer + from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer + + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type == "input": + return HybridQuantizer( + rowwise_quantizer=_plain_nvfp4_quantizer(row_scaled_nvfp4=True), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + if is_linear and role.tensor_type == "weight": + return HybridQuantizer( + rowwise_quantizer=_plain_nvfp4_quantizer(), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + if is_linear and role.tensor_type == "grad_output": + return IdentityQuantizer() + return _plain_nvfp4_quantizer() + + +# ----------------------------------------------------------------------------- +# Linear + Attention Recipes +# ----------------------------------------------------------------------------- + + +def nvfp4_linear_fp8_dpa_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: NVFP4 for ``Linear``, FP8 for ``DotProductAttention``. + + This factory demonstrates how to use ``CustomRecipe`` with ``fp8_dpa=True`` + to combine NVFP4 quantization for linear layers with FP8 attention. + + DPA-owned tensor types (``role.module_type == "dpa"``): + + =========== ============================================================ + tensor_type Description + =========== ============================================================ + ``"qkv"`` Query, Key, Value inputs to the first attention GEMM + ``"s"`` Softmax output (S = softmax(Q·K^T)), fed into the second GEMM + ``"do"`` Gradient of the attention output (dO), backward input + ``"dp"`` Gradient of the softmax output (dP = dO·V^T), backward + =========== ============================================================ + + Dispatch logic: + * ``role.module_type == "dpa"`` with ``tensor_type in ("s", "dp")`` + -> FP8 delayed scaling (``Format.HYBRID``, most_recent, history length 1) + * other DPA roles + -> FP8 current scaling (``Format.HYBRID``: E4M3 fwd, E5M2 bwd) + * DPA boundary hints (``"dpa_output"`` / ``"dpa_grad_input"`` in ``role.name``) + -> FP8 current scaling placeholder. The fused attention kernel requires + FP8-compatible quantizers in all DPA slots, even when the output is + produced in BF16 (``fp8_mha=False``). DPA emits these hint-only roles + (with empty ``module_type`` and ``tensor_type``) when the downstream + consumer is unknown. + * everything else (``"linear"`` / ``"grouped_linear"`` / ``None``) + -> NVFP4 (E2M1), configured per tensor role + + Usage:: + + from transformer_engine.common.recipe import CustomRecipe + from transformer_engine.pytorch.quantization import autocast + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + nvfp4_linear_fp8_dpa_factory, + ) + + recipe = CustomRecipe( + qfactory=nvfp4_linear_fp8_dpa_factory, + fp8_dpa=True, + ) + with autocast(recipe=recipe): + output = model(input) + """ + from transformer_engine.common.recipe import Format + from transformer_engine.pytorch.quantization import DelayedScalingRequest + from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer + + is_dpa = role is not None and role.module_type == "dpa" + is_dpa_boundary = ( + role is not None + and not role.module_type + and ("dpa_output" in role.name or "dpa_grad_input" in role.name) + ) + + # Native NVFP4 + FP8 attention uses delayed scaling for S/dP. + if is_dpa and role.tensor_type in ("s", "dp"): + return DelayedScalingRequest( + fp8_format=Format.HYBRID, + amax_history_len=1, + amax_compute_algo="most_recent", + reduce_amax=True, + ) + + if is_dpa or is_dpa_boundary: + is_bwd_role = (is_dpa and role.tensor_type in ("do", "dp", "dqkv")) or ( + is_dpa_boundary and "dpa_grad_input" in role.name + ) + fp8_dtype = DType.kFloat8E5M2 if is_bwd_role else DType.kFloat8E4M3 + return Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device="cuda") + + return nvfp4_factory(role) diff --git a/transformer_engine/pytorch/custom_recipes/quantization_ref_current_scaling.py b/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py similarity index 97% rename from transformer_engine/pytorch/custom_recipes/quantization_ref_current_scaling.py rename to transformer_engine/pytorch/custom_recipes/reference_current_scaling.py index ecbb667ecf..4e5292019a 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_ref_current_scaling.py +++ b/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py @@ -10,12 +10,12 @@ import torch -from transformer_engine.pytorch.custom_recipes import quantization -from transformer_engine.pytorch.custom_recipes import utils +from transformer_engine.pytorch.custom_recipes import gemm +from transformer_engine.pytorch.custom_recipes import reference_utils from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage, Quantizer -def current_scaling_ref_quantizer_factory(role): +def current_scaling_ref_factory(role): """Factory function for current scaling reference quantizer. Receives a :class:`~transformer_engine.pytorch.quantization.QuantizerRole`. @@ -24,7 +24,7 @@ def current_scaling_ref_quantizer_factory(role): Usage with CustomRecipe and autocast:: - custom_recipe = recipe.CustomRecipe(qfactory=current_scaling_ref_quantizer_factory) + custom_recipe = recipe.CustomRecipe(qfactory=current_scaling_ref_factory) with autocast(recipe=custom_recipe): output = model(input) """ @@ -340,7 +340,9 @@ def quantize( **kwargs, # pylint: disable=unused-argument ) -> CurrentScalingTensorRef: # sanity checks - assert tensor.dtype in utils.HIGH_PRECISION_FLOAT_DTYPES, "Unsupported input dtype." + assert ( + tensor.dtype in reference_utils.HIGH_PRECISION_FLOAT_DTYPES + ), "Unsupported input dtype." # Make it work with 3D tensors original_shape = tensor.shape @@ -374,14 +376,14 @@ def qgemm( self, qx: torch.Tensor, qw: torch.Tensor, - m_params: quantization.MMParams, + m_params: gemm.MMParams, out_dtype: torch.dtype, sx: torch.Tensor, sw: torch.Tensor, bias: torch.Tensor | None = None, out: torch.Tensor | None = None, accumulate: bool = False, - gemm_type: quantization.GEMMType = quantization.GEMMType.FPROP, # pylint: disable=unused-argument + gemm_type: gemm.GEMMType = gemm.GEMMType.FPROP, # pylint: disable=unused-argument qresult_x: QuantizedTensorStorage | None = None, # pylint: disable=unused-argument qresult_w: QuantizedTensorStorage | None = None, # pylint: disable=unused-argument ) -> torch.Tensor: diff --git a/transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py b/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py similarity index 95% rename from transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py rename to transformer_engine/pytorch/custom_recipes/reference_nvfp4.py index fd1fc7544e..3091d4115c 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py +++ b/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py @@ -11,8 +11,8 @@ import torch -from transformer_engine.pytorch.custom_recipes import quantization -from transformer_engine.pytorch.custom_recipes import utils +from transformer_engine.pytorch.custom_recipes import gemm +from transformer_engine.pytorch.custom_recipes import reference_utils from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage, Quantizer from torch.utils.cpp_extension import IS_HIP_EXTENSION @@ -20,7 +20,7 @@ from transformer_engine.pytorch.utils import get_torch_float8_e4m3_type, is_fp8_fnuz -def nvfp4_ref_rht_2d_quantizer_factory(role): +def nvfp4_ref_rht_2d_factory(role): """ Quantizer factory for NVFP4 recipe reference implementation (RHT and 2D quantization for weights). @@ -28,7 +28,7 @@ def nvfp4_ref_rht_2d_quantizer_factory(role): Usage with CustomRecipe and autocast:: - custom_recipe = recipe.CustomRecipe(qfactory=nvfp4_ref_rht_2d_quantizer_factory) + custom_recipe = recipe.CustomRecipe(qfactory=nvfp4_ref_rht_2d_factory) with autocast(recipe=custom_recipe): output = model(input) """ @@ -39,13 +39,13 @@ def nvfp4_ref_rht_2d_quantizer_factory(role): ) if is_weight_tensor_in_gemm: # 2D quantization for weights in GEMM-based modules return NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + dtype=reference_utils.Fp4Formats.E2M1, quant_tile_shape=(16, 16), pow_2_scales=False, with_rht=False, ) return NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + dtype=reference_utils.Fp4Formats.E2M1, quant_tile_shape=(1, 16), pow_2_scales=False, with_rht=True, @@ -219,7 +219,7 @@ class NVFP4TensorRef(QuantizedTensorStorage): nominal tensor datatype. device: torch.device device of the tensor. - quant_dtype: Union[utils.Fp4Formats, torch.dtype] + quant_dtype: Union[reference_utils.Fp4Formats, torch.dtype] low precision tensor datatype. original_shape: Tuple[int, ...] original shape of the tensor. @@ -238,7 +238,7 @@ class NVFP4TensorRef(QuantizedTensorStorage): dtype: Optional[torch.dtype] = None device: Optional[torch.device] = None - quant_dtype: Optional[Union[utils.Fp4Formats, torch.dtype]] = None + quant_dtype: Optional[Union[reference_utils.Fp4Formats, torch.dtype]] = None original_shape: Optional[Tuple[int, ...]] = None _quantizer: Optional[Quantizer] = None @@ -357,7 +357,7 @@ class NVFP4QuantizerRef(Quantizer): def __init__( self, - dtype: utils.Fp4Formats, + dtype: reference_utils.Fp4Formats, rowwise: bool = True, columnwise: bool = True, pow_2_scales: bool = False, @@ -375,10 +375,6 @@ def __init__( if row_scaled_nvfp4: if not rowwise: raise ValueError("Row-scaled NVFP4 reference quantization requires rowwise usage.") - if columnwise: - raise ValueError( - "Row-scaled NVFP4 reference quantization does not support columnwise usage." - ) if nvfp4_use_4over6: if nvfp4_4over6_err_mode not in ("MAE", "MSE"): raise ValueError(f"Unsupported NVFP4 4over6 error mode: {nvfp4_4over6_err_mode}.") @@ -468,9 +464,9 @@ def _recover_swizzled_scales( ) -> torch.Tensor: if not swizzled_scale: return scale - rounded_m = utils.roundup_div(m, 128) * 128 - scale_n = utils.roundup_div(n, block_length) - rounded_n = utils.roundup_div(scale_n, 4) * 4 + rounded_m = reference_utils.roundup_div(m, 128) * 128 + scale_n = reference_utils.roundup_div(n, block_length) + rounded_n = reference_utils.roundup_div(scale_n, 4) * 4 # Recover swizzled scaling factor layout -> linear layout tmp = torch.reshape(scale, (rounded_m // 128, rounded_n // 4, 32, 4, 4)) # after permutation, the layout is [rounded_m // 128, 4, 32, rounded_n // 4, 4] @@ -901,7 +897,14 @@ def _quantize(self, tensor: torch.Tensor) -> Tuple[ f"got {self.quant_tile_shape}" ) global_amax_row = torch.max(torch.abs(row_input), dim=1).values.to(torch.float32) - global_amax_col = global_amax_row + # Columnwise (transpose) uses per-row-of-transpose amax, i.e. the + # per-column amax of the original input. When columnwise output is + # not requested, keep it aliased to the rowwise amax as before. + global_amax_col = ( + torch.max(torch.abs(col_input), dim=1).values.to(torch.float32) + if self.columnwise_usage + else global_amax_row + ) else: # Compute amax for rowwise and columnwise paths separately global_amax_row = torch.max(torch.abs(row_input)).to(torch.float32).view(1) @@ -954,6 +957,7 @@ def _quantize(self, tensor: torch.Tensor) -> Tuple[ self.quant_tile_shape[1], self.quant_tile_shape[0], pow_2_scales=self.pow_2_scales, + row_scaled_nvfp4=self.row_scaled_nvfp4, nvfp4_use_4over6=self.nvfp4_use_4over6, nvfp4_e4m3_max=self.nvfp4_e4m3_max, nvfp4_4over6_err_mode=self.nvfp4_4over6_err_mode, @@ -977,10 +981,10 @@ def quantize( **kwargs, # pylint: disable=unused-argument ) -> NVFP4TensorRef: # sanity checks - if tensor.dtype not in utils.HIGH_PRECISION_FLOAT_DTYPES: + if tensor.dtype not in reference_utils.HIGH_PRECISION_FLOAT_DTYPES: raise TypeError( f"Unsupported input dtype {tensor.dtype}, expected one of" - f" {utils.HIGH_PRECISION_FLOAT_DTYPES}" + f" {reference_utils.HIGH_PRECISION_FLOAT_DTYPES}" ) # Make it work with 3D tensors @@ -1090,14 +1094,14 @@ def qgemm( self, qx: torch.Tensor, qw: torch.Tensor, - m_params: quantization.MMParams, # pylint: disable=unused-argument + m_params: gemm.MMParams, # pylint: disable=unused-argument out_dtype: torch.dtype, sx: torch.Tensor, sw: torch.Tensor, bias: torch.Tensor | None = None, out: torch.Tensor | None = None, accumulate: bool = False, - gemm_type: quantization.GEMMType = quantization.GEMMType.FPROP, + gemm_type: gemm.GEMMType = gemm.GEMMType.FPROP, qresult_x: QuantizedTensorStorage | None = None, qresult_w: QuantizedTensorStorage | None = None, ) -> torch.Tensor: @@ -1184,14 +1188,23 @@ def qgemm( fp8_max_w = default_fp8_max factor = 6.0 * 6.0 * fp8_max_x * fp8_max_w - if gemm_type == quantization.GEMMType.WGRAD: + if gemm_type == gemm.GEMMType.WGRAD: partial_alpha = qresult_x.global_amax_col * qresult_w.global_amax_col + # A row-scaled operand contributes a per-output-column (N) vector + # here, so broadcast along the last axis. Selecting the axis from + # gemm_type (rather than matching numel against M) avoids the + # square-matrix ambiguity where M == N. + if partial_alpha.numel() > 1: + partial_alpha = partial_alpha.reshape(1, -1) + else: + partial_alpha = partial_alpha.squeeze(-1) else: partial_alpha = qresult_x.global_amax_row * qresult_w.global_amax_row - if partial_alpha.numel() > 1 and partial_alpha.numel() == high_precision_x.shape[0]: - partial_alpha = partial_alpha.view(-1, 1) - else: - partial_alpha = partial_alpha.squeeze(-1) + # A row-scaled operand contributes a per-output-row (M) vector. + if partial_alpha.numel() > 1: + partial_alpha = partial_alpha.reshape(-1, 1) + else: + partial_alpha = partial_alpha.squeeze(-1) alpha = torch.div(partial_alpha, factor) M, K = high_precision_x.shape diff --git a/transformer_engine/pytorch/custom_recipes/utils.py b/transformer_engine/pytorch/custom_recipes/reference_utils.py similarity index 83% rename from transformer_engine/pytorch/custom_recipes/utils.py rename to transformer_engine/pytorch/custom_recipes/reference_utils.py index 3e23661f14..5bd25e35b2 100644 --- a/transformer_engine/pytorch/custom_recipes/utils.py +++ b/transformer_engine/pytorch/custom_recipes/reference_utils.py @@ -2,7 +2,7 @@ # # See LICENSE for license information. -"""Utility functions for experimental middleware between Transformer Engine and Kitchen.""" +"""Shared utilities for custom recipe reference implementations.""" import enum diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index ffb641a38c..a273f2a0df 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -5,10 +5,11 @@ # See LICENSE for license information. """Methods needed for distributed training (DP/TP).""" + from __future__ import annotations from collections.abc import Iterable -from contextlib import contextmanager, AbstractContextManager, ContextDecorator +from contextlib import contextmanager, AbstractContextManager, ContextDecorator, nullcontext from functools import lru_cache from dataclasses import dataclass import math @@ -52,7 +53,6 @@ from .tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from ..debug.pytorch.debug_quantization import DebugQuantizedTensor - __all__ = ["checkpoint", "CudaRNGStatesTracker"] @@ -64,8 +64,8 @@ _USE_REENTRANT_ACTIVATION_RECOMPUTE = True -_FP8_ACTIVATION_RECOMPUTE_ENABLED = False -_FP8_ACTIVATION_RECOMPUTE_PHASE = False +_IN_ACTIVATION_RECOMPUTE_REGION = False +_ACTIVATION_RECOMPUTE_PHASE = False _ALL_ACTIVE_RNG_STATES = {} @@ -257,11 +257,14 @@ def __init__(self, activation_recompute: bool = False, recompute_phase: bool = F self.recompute_phase = recompute_phase def __enter__(self): - global _FP8_ACTIVATION_RECOMPUTE_ENABLED, _FP8_ACTIVATION_RECOMPUTE_PHASE - _FP8_ACTIVATION_RECOMPUTE_ENABLED = ( - self.activation_recompute and FP8GlobalStateManager.is_fp8_enabled() - ) - _FP8_ACTIVATION_RECOMPUTE_PHASE = self.recompute_phase + global _IN_ACTIVATION_RECOMPUTE_REGION, _ACTIVATION_RECOMPUTE_PHASE + # Track the checkpoint region independently of the FP8 state at entry. + # A checkpointed callable may open its own FP8 autocast context (for + # example, to select precision per layer). Delayed-scaling modules in + # that inner context must still save their scale and amax metadata for + # the recompute forward. + _IN_ACTIVATION_RECOMPUTE_REGION = self.activation_recompute + _ACTIVATION_RECOMPUTE_PHASE = self.recompute_phase qstate = FP8GlobalStateManager.quantization_state if self.activation_recompute and not self.recompute_phase: @@ -270,19 +273,19 @@ def __enter__(self): qstate.is_first_fp8_module = activation_recompute_forward._is_first_fp8_module.pop(0) def __exit__(self, *exc_details): - global _FP8_ACTIVATION_RECOMPUTE_ENABLED, _FP8_ACTIVATION_RECOMPUTE_PHASE - _FP8_ACTIVATION_RECOMPUTE_ENABLED = False - _FP8_ACTIVATION_RECOMPUTE_PHASE = False + global _IN_ACTIVATION_RECOMPUTE_REGION, _ACTIVATION_RECOMPUTE_PHASE + _IN_ACTIVATION_RECOMPUTE_REGION = False + _ACTIVATION_RECOMPUTE_PHASE = False def is_fp8_activation_recompute_enabled() -> bool: - """Return global boolean""" - return _FP8_ACTIVATION_RECOMPUTE_ENABLED + """Whether we are in an activation recompute region with FP8 currently enabled""" + return _IN_ACTIVATION_RECOMPUTE_REGION and FP8GlobalStateManager.is_fp8_enabled() def in_fp8_activation_recompute_phase() -> bool: """Return global boolean""" - return _FP8_ACTIVATION_RECOMPUTE_PHASE + return _ACTIVATION_RECOMPUTE_PHASE def _get_active_autocast_contexts(): @@ -930,7 +933,10 @@ def fork(self, name: str = "model-parallel-rng"): def reduce_scatter_along_first_dim( - inp: torch.Tensor, tp_group: dist_group_type, async_op: bool = False + inp: torch.Tensor, + tp_group: dist_group_type, + async_op: bool = False, + output: torch.Tensor = None, ) -> Tuple[torch.Tensor, Optional[torch.distributed.Work]]: """Reduce-scatter the input tensor across model parallel group.""" world_size = get_distributed_world_size(tp_group) @@ -948,7 +954,8 @@ def reduce_scatter_along_first_dim( dim_size[0] = dim_size[0] // world_size - output = torch.empty(dim_size, dtype=inp.dtype, device=torch.cuda.current_device()) + if output is None: + output = torch.empty(dim_size, dtype=inp.dtype, device=torch.cuda.current_device()) handle = torch.distributed.reduce_scatter_tensor( output, inp.contiguous(), group=tp_group, async_op=async_op ) @@ -1315,7 +1322,8 @@ def wait(self) -> None: """Wait for the async operation to complete and post-process the tensor.""" if self._synchronized: return - self.async_handle.wait() + if self.async_handle is not None: + self.async_handle.wait() _post_process_nvfp4_gather( self.output, self.columnwise_data_interleaved, @@ -1332,6 +1340,8 @@ def _all_gather_nvfp4( async_op: bool = False, quantizer: NVFP4Quantizer, out_shape: Optional[list[int]] = None, + output_tensor=None, + external_coalescing=False, ) -> tuple[NVFP4TensorStorage, Optional[torch.distributed.Work]]: """All-gather NVFP4 tensor along first dimension.""" @@ -1408,15 +1418,23 @@ def _all_gather_nvfp4( inp = quantizer(inp.dequantize(dtype=dtype)) # Construct NVFP4 output tensor - out = quantizer.make_empty(out_shape, dtype=dtype, device=device) + if output_tensor is not None: + out = output_tensor + else: + out = quantizer.make_empty(out_shape, dtype=dtype, device=device) # Coalesce NCCL collectives for gathering data and scale inverses. - with torch.distributed._coalescing_manager( - group=process_group, - device=device, - async_ops=async_op, - ) as gather_coalescing_manager: + if not external_coalescing: + gather_coalescing_manager = torch.distributed._coalescing_manager( + group=process_group, + device=device, + async_ops=async_op, + ) + else: + # Caller owns an outer coalescing manager (managers cannot nest); step aside. + gather_coalescing_manager = nullcontext() + with gather_coalescing_manager as coalesced_handle: # Gather NVFP4 data for row-wise usage if quantizer.rowwise_usage: @@ -1497,10 +1515,10 @@ def _all_gather_nvfp4( # Transfer amax to output. out._amax_columnwise = inp._amax_columnwise - handle = gather_coalescing_manager if async_op else None + handle = coalesced_handle if async_op else None # Fixes interleaved data for transposed tensor/scale inv and pads scale inv if needed. - if async_op and quantizer.columnwise_usage: + if (async_op or external_coalescing) and quantizer.columnwise_usage: handle = _NVFP4AllGatherAsyncHandle( out, out_columnwise_data, out_scale_inv, world_size, handle ) @@ -1517,6 +1535,8 @@ def _all_gather_mxfp8( async_op: bool = False, quantizer: MXFP8Quantizer, out_shape: Optional[list[int]] = None, + output_tensor: torch.Tensor = None, + external_coalescing: bool = False, ) -> tuple[MXFP8TensorStorage, Optional[torch.distributed.Work]]: """All-gather MXFP8 tensor along first dimension.""" @@ -1582,15 +1602,23 @@ def _all_gather_mxfp8( inp = quantizer(inp.dequantize(dtype=dtype)) # Construct MXFP8 output tensor - out = quantizer.make_empty(out_shape, dtype=dtype, device=device) + if output_tensor is not None: + out = output_tensor + else: + out = quantizer.make_empty(out_shape, dtype=dtype, device=device) - # Coalesce NCCL collectives - with torch.distributed._coalescing_manager( - group=process_group, - device=device, - async_ops=async_op, - ) as coalescing_manager: + if not external_coalescing: + # Coalesce NCCL collectives for gathering data and scale inverses. + gather_coalescing_manager = torch.distributed._coalescing_manager( + group=process_group, + device=device, + async_ops=async_op, + ) + else: + # Caller owns an outer coalescing manager (managers cannot nest); step aside. + gather_coalescing_manager = nullcontext() + with gather_coalescing_manager as coalesced_handle: # Gather MXFP8 data for row-wise usage if quantizer.rowwise_usage: @@ -1637,7 +1665,7 @@ def _all_gather_mxfp8( group=process_group, ) - handle = coalescing_manager if async_op else None + handle = coalesced_handle if async_op else None return out, handle @@ -1646,9 +1674,17 @@ def gather_along_first_dim( process_group: dist_group_type, async_op: bool = False, quantizer: Optional[Quantizer] = None, + output_tensor: torch.Tensor = None, + external_coalescing: bool = False, ) -> tuple[torch.Tensor, Optional[torch.distributed.Work]]: """ All-gather tensors and concatenate along first dimension. + + ``external_coalescing``: composability flag for callers that batch several gathers into + one outer ``torch.distributed._coalescing_manager``. Coalescing managers cannot nest, so + when set this call skips opening its own manager and defers any post-gather fixup + (e.g. NVFP4 columnwise de-interleave) into the returned handle; ``handle.wait()`` completes + it once the outer manager has closed. Leave ``False`` for standalone gathers. """ # Return immediately if no communication is required @@ -1736,6 +1772,8 @@ def gather_along_first_dim( async_op=async_op, quantizer=quantizer, out_shape=out_shape, + output_tensor=output_tensor, + external_coalescing=external_coalescing, ) # NVFP4 case @@ -1750,6 +1788,8 @@ def gather_along_first_dim( async_op=async_op, quantizer=quantizer, out_shape=out_shape, + output_tensor=output_tensor, + external_coalescing=external_coalescing, ) # High-precision communication for quantized tensors @@ -1779,19 +1819,20 @@ def gather_along_first_dim( inp = inp.dequantize() # Communication for plain PyTorch tensors - out = torch.empty( - out_shape, - dtype=inp.dtype, - device=inp.device, - memory_format=torch.contiguous_format, - ) + if output_tensor is None: + output_tensor = torch.empty( + out_shape, + dtype=inp.dtype, + device=inp.device, + memory_format=torch.contiguous_format, + ) handle = torch.distributed.all_gather_into_tensor( - out, + output_tensor, inp.contiguous(), group=process_group, async_op=async_op, ) - return out, handle + return output_tensor, handle # Global cache to store symmetric memory tensors @@ -1846,13 +1887,95 @@ def get_symmetric_memory_tensor(tensor_numel, tensor_dtype, tensor_device, tp_gr return msg +_SYMM_MEM_POOL = None +_SYMM_MEM_POOL_BACKEND = None +# Device the pool was created for; the torch symm_mem._symm_mem_pools cache is keyed by it. +_SYMM_MEM_POOL_DEVICE = None +# True when the pool was created via torch's get_mem_pool, which caches it in the private +# symm_mem._symm_mem_pools dict; release then has to drop that cached reference. +_SYMM_MEM_POOL_TORCH_CACHED = False + + +def _get_symm_mem_pool(device: torch.device, backend: str = "NCCL"): + """Process-wide torch MemPool backed by the symmetric-memory allocator, created once (each rank + drives one device). The pool/allocator captures the backend at creation and there is no per-pool + backend arg, so the (process-global) backend is always set before the pool is created. The + collective rendezvous cost is amortized across allocations (paid per new segment, not per buffer). + """ + global _SYMM_MEM_POOL, _SYMM_MEM_POOL_BACKEND, _SYMM_MEM_POOL_DEVICE, _SYMM_MEM_POOL_TORCH_CACHED + if _SYMM_MEM_POOL is None: + symm_mem.set_backend(backend) + _SYMM_MEM_POOL_BACKEND = backend + _SYMM_MEM_POOL_DEVICE = device + if hasattr(symm_mem, "get_mem_pool"): + _SYMM_MEM_POOL = symm_mem.get_mem_pool(device) + _SYMM_MEM_POOL_TORCH_CACHED = True + elif hasattr(torch.cuda, "MemPool") and hasattr(symm_mem, "get_mempool_allocator"): + _SYMM_MEM_POOL = torch.cuda.MemPool(symm_mem.get_mempool_allocator(device)) + else: + raise RuntimeError( + "No symmetric-memory MemPool API available (need torch symm-mem get_mem_pool, or " + "torch.cuda.MemPool + get_mempool_allocator)." + ) + elif backend != _SYMM_MEM_POOL_BACKEND: + raise RuntimeError( + f"symm-mem pool already created with backend {_SYMM_MEM_POOL_BACKEND!r}; " + f"cannot switch to {backend!r}" + ) + return _SYMM_MEM_POOL + + +def release_symm_mem_pool() -> None: + """Free the process-wide symm-mem pool's segments, deregistering their NCCL windows. + + Call before ``dist.destroy_process_group()``: the pool's windows are registered on + the group's NCCL comm, which becomes invalid once the group is destroyed. No-op if + no pool was created. + """ + global _SYMM_MEM_POOL, _SYMM_MEM_POOL_BACKEND, _SYMM_MEM_POOL_DEVICE, _SYMM_MEM_POOL_TORCH_CACHED + if _SYMM_MEM_POOL is None: + return + # The torch symm_mem._symm_mem_pools cache is keyed by the pool's creation device. + device = _SYMM_MEM_POOL_DEVICE + _SYMM_MEM_POOL = None + _SYMM_MEM_POOL_BACKEND = None + _SYMM_MEM_POOL_DEVICE = None + torch_cached = _SYMM_MEM_POOL_TORCH_CACHED + _SYMM_MEM_POOL_TORCH_CACHED = False + # A pool from torch's get_mem_pool is also cached in the private module dict + # symm_mem._symm_mem_pools; drop that reference so the segments' refcount reaches zero + # and their NCCL windows deregister. Fail loudly if this internal has changed shape, + # otherwise the windows would silently leak past destroy_process_group(). + if torch_cached: + pools = getattr(symm_mem, "_symm_mem_pools", None) + if not isinstance(pools, dict) or device not in pools: + raise RuntimeError( + "torch symmetric-memory pool cache (symm_mem._symm_mem_pools) is missing or " + "has changed layout; cannot release the pooled segments and their NCCL windows " + "would leak past destroy_process_group(). This torch version needs an updated " + "release_symm_mem_pool()." + ) + pools.pop(device, None) + torch.cuda.empty_cache() + + def symm_mem_alloc( shape, dtype: torch.dtype, ep_group: dist_group_type, device: Optional[torch.device] = None, + use_pool: bool = False, + backend: str = "NCCL", ) -> torch.Tensor: - """Allocate and rendezvous a symm-mem buffer on ep_group. Collective on ep_group.""" + """Allocate a symm-mem buffer on ep_group. + + ``use_pool=False`` (default): freshly allocate and do one explicit collective ``rendezvous`` per + buffer, for caller-owned static buffers (fp8 zero-copy). ``use_pool=True``: allocate from a + process-wide symm-mem MemPool whose segments are auto-registered (implicit mempool), so no explicit + rendezvous is needed and torch manages the tensor lifecycle (freed back to the pool) — for + lifecycle-managed zero-copy, e.g. bf16, where the recv buffer is saved for backward and so cannot + be a shared static buffer. ``backend`` selects the symm-mem backend (default NCCL; for the pool it + is captured at pool creation).""" if device is None: device = torch.device("cuda", torch.cuda.current_device()) if not HAS_TORCH_SYMMETRIC: @@ -1860,10 +1983,15 @@ def symm_mem_alloc( "torch.distributed._symmetric_memory is unavailable; symm_mem_alloc " "requires PyTorch built with NCCL symm-mem support." ) - if symm_mem.get_backend(device) != "NCCL": - symm_mem.set_backend("NCCL") - t = symm_mem.empty(*shape, dtype=dtype, device=device) - symm_mem.rendezvous(t, group=ep_group) + if use_pool: + pool = _get_symm_mem_pool(device, backend) + with torch.cuda.use_mem_pool(pool): + t = torch.empty(*shape, dtype=dtype, device=device) + else: + if symm_mem.get_backend(device) != backend: + symm_mem.set_backend(backend) + t = symm_mem.empty(*shape, dtype=dtype, device=device) + symm_mem.rendezvous(t, group=ep_group) return t diff --git a/transformer_engine/pytorch/distributed_weight.py b/transformer_engine/pytorch/distributed_weight.py new file mode 100644 index 0000000000..0d491430e5 --- /dev/null +++ b/transformer_engine/pytorch/distributed_weight.py @@ -0,0 +1,117 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""GTP-agnostic weight-parallelism extension point. + +TE owns this contract but ships no implementation; the caller (e.g. Megatron GTP) implements the +protocol on the weight and injects it at construction. Dispatchers are list-shaped (Linear -> 1, +GroupedLinear -> N; leader is ``weights[0]``) and no-op on plain tensors. +""" + +from typing import Any, List, Protocol, runtime_checkable + +import torch + +__all__ = [ + "DistributedWeight", + "is_distributed_weight", + "materialize_weight_for_forward", + "materialize_weight_for_backward", + "finalize_weight_grads", +] + + +@runtime_checkable +class DistributedWeight(Protocol): + """Structural interface for a custom-weight-parallel weight (AG for the GEMM, reduce/RS the + grad, re-materialize in backward). Duck-typed ``typing.Protocol``: implementers need not + subclass it, and all state (shards, group, async handles) lives outside TE on the implementer. + + Implementers MUST be ``torch.Tensor`` subclasses (needed by ``ctx.save_for_backward``, DDP + backward hooks, and ``torch.compile``); enforced at runtime by :func:`is_distributed_weight`. + """ + + # Capability marker: True on an implementer, absent on plain tensors; TE's fwd/bwd gate on it. + is_distributed_weight: bool + + def materialize_group_for_forward(self) -> Any: + """Return the tensor(s) to feed the forward GEMM (may all-gather shards).""" + + def materialize_group_for_backward(self) -> Any: + """Re-materialize the full weight(s) for the backward GEMMs.""" + + def finalize_group_grads(self, wgrads: Any) -> Any: + """Post-process freshly computed weight grad(s) (e.g. reduce-scatter). + + May consume ``wgrads`` in-place -- reduce-scatter into ``main_grad`` and set + ``grad_added_to_main_grad`` -- returning a dummy grad (or ``None`` for an async collective) + that callers use as the parameter grad(s) or discard. + """ + + def grad_buffer(self) -> torch.Tensor: + """The gradient accumulation buffer for this weight.""" + + +def is_distributed_weight(weight: Any) -> bool: + """True if ``weight`` participates in custom weight parallelism (False on plain tensors). + + Enforces the :class:`DistributedWeight` requirement that an implementer be a ``torch.Tensor`` + subclass, failing loudly here rather than silently breaking autograd downstream. + """ + flag = bool(getattr(weight, "is_distributed_weight", False)) + if flag and not isinstance(weight, torch.Tensor): + raise TypeError( + "DistributedWeight implementers must be torch.Tensor subclasses; got " + f"{type(weight).__name__}." + ) + return flag + + +def materialize_weight_for_forward(weights: Any) -> List[Any]: + """Prepare the weight(s) fed to the forward GEMM, always returned as a list. + + Args: + weights: the module's weight(s) -- a single weight (Linear) or the full per-expert list + (GroupedLinear). A bare weight is treated as a one-element list. + + Returns: + - Distributed group: the leader ``weights[0]`` all-gathers/coalesces the whole group and + returns all N materialized weights; the follower entries ``weights[1:]`` are ignored + here (the leader already holds references to its group). + - Otherwise: the input weights, unchanged. + """ + if not isinstance(weights, (list, tuple)): + weights = [weights] + leader = weights[0] + if is_distributed_weight(leader): + out = leader.materialize_group_for_forward() + return list(out) if isinstance(out, (list, tuple)) else [out] + return list(weights) + + +def materialize_weight_for_backward(weights: Any) -> List[Any]: + """Backward-GEMM mirror of :func:`materialize_weight_for_forward` (same contract).""" + if not isinstance(weights, (list, tuple)): + weights = [weights] + leader = weights[0] + if is_distributed_weight(leader): + out = leader.materialize_group_for_backward() + return list(out) if isinstance(out, (list, tuple)) else [out] + return list(weights) + + +def finalize_weight_grads(weights: Any, wgrads: List[Any]) -> List[Any]: + """Finalize a weight group's grad(s), mirroring :func:`materialize_weight_for_backward`. + + Delegates to the leader's :meth:`DistributedWeight.finalize_group_grads` (which defines the + in-place / dummy / async-``None`` return contract); returns ``wgrads`` unchanged when not + distributed. + """ + if not isinstance(weights, (list, tuple)): + weights = [weights] + leader = weights[0] + if is_distributed_weight(leader): + out = leader.finalize_group_grads(wgrads if len(wgrads) > 1 else wgrads[0]) + return list(out) if isinstance(out, (list, tuple)) else [out] + return list(wgrads) diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py index ee860c78e3..4d5c76e9ce 100644 --- a/transformer_engine/pytorch/dynamo/__init__.py +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -5,8 +5,11 @@ """torch.compile glue for Transformer Engine.""" from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer +from .tensor_spec import TensorSpec, to_tensor_spec __all__ = [ "register_value_opaque_quantizer", "is_value_opaque_quantizer", + "TensorSpec", + "to_tensor_spec", ] diff --git a/transformer_engine/pytorch/dynamo/tensor_spec.py b/transformer_engine/pytorch/dynamo/tensor_spec.py new file mode 100644 index 0000000000..4cfe225952 --- /dev/null +++ b/transformer_engine/pytorch/dynamo/tensor_spec.py @@ -0,0 +1,165 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""TensorSpec: a data-free description of a tensor / quantized tensor.""" + +from __future__ import annotations +import copy as _copy +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Tuple + +import torch +from torch._prims_common import make_contiguous_strides_for + + +@dataclass +class TensorSpec: + """A data-free description of a tensor or quantized tensor. + + Captures ``shape`` / ``dtype`` and, for quantized tensors, the + (value-opaque) ``quantizer`` -- enough to rebuild a tensor without holding + storage. The common abstraction over plain ``torch.Tensor``, + ``QuantizedTensorStorage`` and ``QuantizedTensor``, used for custom-op fake + impls and for reassembling a quantized tensor from bare inner tensors. + """ + + shape: Tuple[int, ...] + dtype: torch.dtype + quantizer: Optional[Any] = None + requires_grad: bool = False + device: Optional[torch.device] = field(default=None) + + def __post_init__(self) -> None: + # Own a private copy of the quantizer so usage changes (update_usage) + # never touch the shared, value-opaque quantizer. The copy inherits the + # quantizer's current row-/column-wise usage as this spec's layout. + if self.quantizer is not None: + q = self.quantizer + self.quantizer = q.copy() if hasattr(q, "copy") else _copy.copy(q) + + @property + def is_quantized(self) -> bool: + """Whether this spec describes a quantized tensor.""" + return self.quantizer is not None + + def update_usage( + self, + *, + rowwise_usage: Optional[bool] = None, + columnwise_usage: Optional[bool] = None, + ) -> None: + """Mirror ``QuantizedTensor.update_usage`` on the spec's inner-tensor layout. + + Applied to the spec's own quantizer copy, so the shared (value-opaque) + quantizer is never mutated. Raises on plain (non-quantized) specs -- + a real plain ``torch.Tensor`` has no ``update_usage`` either. + """ + if self.quantizer is None: + raise ValueError("update_usage called on a non-quantized TensorSpec") + self.quantizer.set_usage(rowwise=rowwise_usage, columnwise=columnwise_usage) + + def inner_names(self) -> Tuple[str, ...]: + """Names of the flat inner tensors backing this spec, in order. + + The real op flattens a quantized output via the storage's + ``__tensor_flatten__`` -- i.e. ``_INNER_TENSORS`` order, keeping only the + present inner tensors. ``inner_tensor_specs`` is contracted to emit them + in that same order, which keeps the fake layout aligned with the real one + slot-for-slot; the contract is checked here rather than papered over by + reordering, so a mismatching quantizer fails loudly. + """ + if self.quantizer is None: + return ("data",) + # pylint: disable=protected-access + described = tuple(self.quantizer.inner_tensor_specs(tuple(self.shape))) + storage_cls = self.quantizer.storage_metadata(self.dtype)["cls"] + flatten_order = tuple(attr for attr, _ in storage_cls._INNER_TENSORS) + expected = tuple(name for name in flatten_order if name in described) + if described != expected: + raise RuntimeError( + f"{type(self.quantizer).__name__}.inner_tensor_specs returned {described}, " + f"which does not follow {storage_cls.__name__}._INNER_TENSORS order " + f"{flatten_order} (expected {expected}); the fake layout would not match " + "the real one slot-for-slot." + ) + return described + + def create_metadata(self) -> Dict[str, Any]: + """Data-free ``__tensor_unflatten__`` context describing this tensor.""" + if self.quantizer is None: + return { + "is_tensor": True, + "is_quantized": False, + "dtype": self.dtype, + "requires_grad": self.requires_grad, + } + return self.quantizer.create_metadata( + tuple(self.shape), dtype=self.dtype, requires_grad=self.requires_grad + ) + + def create_inner_tensors(self) -> List[torch.Tensor]: + """Materialize the flat inner tensors (in :meth:`inner_names` order). + + Under ``register_fake`` the ``torch.empty`` calls produce ``FakeTensor``s; + ``requires_grad`` is left default (managed by ``register_autograd``). + """ + device = self.device if self.device is not None else torch.device("cuda") + if self.quantizer is None: + return [torch.empty(tuple(self.shape), dtype=self.dtype, device=device)] + inner = self.quantizer.alloc_tensors(tuple(self.shape), device=device) + return [inner[name] for name in self.inner_names()] + + def assemble(self, inner_tensors: List[torch.Tensor]) -> torch.Tensor: + """Rebuild the tensor from ready-made ``inner_tensors`` (in :meth:`inner_names` + order). Shared by :meth:`create_tensor` (fresh ones) and the custom-op + boundary (inner tensors arriving from an op's flat ``Tensor[]`` payload). + + Non-quantized specs are the single inner tensor as-is; quantized specs + are reassembled into the storage/wrapper via ``__tensor_unflatten__``. + """ + if self.quantizer is None: + return inner_tensors[0] + shape = tuple(self.shape) + ctx = self.create_metadata() + inner = dict(zip(self.inner_names(), inner_tensors)) + storage_cls = ctx["cls"] + return storage_cls.__tensor_unflatten__( + inner, ctx, shape, make_contiguous_strides_for(shape) + ) + + def create_tensor(self) -> torch.Tensor: + """Materialize an (uninitialized) tensor matching this spec (traceable). + + Quantized specs reassemble freshly-allocated :meth:`create_inner_tensors` + inner tensors via :meth:`assemble`. + """ + if self.quantizer is None: + device = self.device if self.device is not None else torch.device("cuda") + return torch.empty( + tuple(self.shape), + dtype=self.dtype, + device=device, + requires_grad=self.requires_grad, + ) + return self.assemble(self.create_inner_tensors()) + + +def to_tensor_spec(tensor: Any) -> TensorSpec: + """Build a :class:`TensorSpec` describing ``tensor``. + + Works for plain ``torch.Tensor`` and for ``QuantizedTensorStorage`` / + ``QuantizedTensor``. A *bare* storage exposes its (fake) dtype via + ``_dtype`` rather than ``.dtype``. + """ + requires_grad = bool(getattr(tensor, "requires_grad", False)) + dtype = getattr(tensor, "dtype", None) + if dtype is None: + dtype = getattr(tensor, "_dtype", None) + return TensorSpec( + shape=tuple(tensor.shape), + dtype=dtype, + quantizer=getattr(tensor, "_quantizer", None), + requires_grad=requires_grad, + device=tensor.device, + ) diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index e57caa1f42..2799bfdf5d 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -7,7 +7,7 @@ import atexit import warnings -from typing import Optional +from typing import Optional, TYPE_CHECKING import torch import torch.distributed as dist @@ -15,16 +15,24 @@ import transformer_engine_torch as tex from .cpu_offload import mark_not_offload -from .distributed import symm_mem_alloc +from .distributed import symm_mem_alloc, release_symm_mem_pool +from .quantized_tensor import QuantizedTensor +# Type-hint-only import; keeps the ``Recipe`` annotation without a runtime import of +# common.recipe (the concrete recipe classes are imported lazily where used). +if TYPE_CHECKING: + from ..common.recipe import Recipe __all__ = [ "EpBuffer", "ep_bootstrap", + "is_ep_bootstrapped", "ep_finalize", "ep_dispatch", "ep_combine", "symm_mem_alloc", + "release_symm_mem_pool", + "is_symm_backed", ] @@ -66,14 +74,17 @@ def _check_nccl_runtime_version() -> None: _BOOTSTRAPPED = False _ATEXIT_REGISTERED = False -# EP group captured at bootstrap; EpBuffer uses it to allocate the symm-mem -# combine grad buffer in zero-copy mode. +# EP group captured at bootstrap; used by the zero-copy symm-mem pool allocator. _EP_GROUP: Optional[dist.ProcessGroup] = None +# Eager-mode toggle captured at bootstrap (set when recv_capacity_per_rank is +# omitted); ep_dispatch reads it to size the recv outputs from the per-step +# recv-token total instead of a fixed recv_capacity_per_rank. +_EAGER = False def _atexit_finalize() -> None: """Best-effort teardown at interpreter shutdown; swallows errors.""" - global _BOOTSTRAPPED, _EP_GROUP + global _BOOTSTRAPPED, _EP_GROUP, _EAGER if _BOOTSTRAPPED: try: tex.ep_finalize() @@ -84,16 +95,19 @@ def _atexit_finalize() -> None: finally: _BOOTSTRAPPED = False _EP_GROUP = None + _EAGER = False def ep_bootstrap( ep_group: dist.ProcessGroup, num_experts: int, max_tokens_per_rank: int, - recv_capacity_per_rank: int, hidden_dim: int, + num_topk: int, + recv_capacity_per_rank: Optional[int] = None, max_num_sms: int = 0, zero_copy: bool = False, + drop_on_overflow: bool = False, max_token_dtype: torch.dtype = torch.bfloat16, ) -> None: """Initialize EP by borrowing ep_group's NCCL comm. Call once per process. @@ -101,15 +115,34 @@ def ep_bootstrap( max_token_dtype sets the widest token dtype this EP group will dispatch; it sizes NCCL EP staging buffers. + ``recv_capacity_per_rank`` bounds the tokens one rank receives per step and + sizes the recv outputs. Omit it (``None``) for eager mode, which sizes recv + outputs from the per-step recv total instead; eager needs a host sync each + step and is not CUDA-graph capturable. + ``zero_copy`` opts the EP group into the symm-mem zero-copy IO path; pass ``True`` only when payload tensors are allocated via ``symm_mem_alloc``. - Defaults to ``False``. + Requires ``recv_capacity_per_rank``. To capture a CUDA graph, supply + persistent recv_tokens / grad_out buffers to dispatch/combine; the pool-based + auto-allocation used when they are omitted is not CUDA-graph capturable. + + ``num_topk`` is the per-token top-k; it sizes NCCL EP internal buffers. + + ``drop_on_overflow`` drops tokens exceeding ``recv_capacity_per_rank`` instead + of trapping. Requires ``recv_capacity_per_rank``. """ - global _BOOTSTRAPPED, _ATEXIT_REGISTERED, _EP_GROUP + global _BOOTSTRAPPED, _ATEXIT_REGISTERED, _EP_GROUP, _EAGER + eager = recv_capacity_per_rank is None if _BOOTSTRAPPED: raise RuntimeError("ep_bootstrap was already called in this process") if ep_group.size() < 2: raise ValueError(f"ep_bootstrap requires ep_group.size() >= 2 (got {ep_group.size()}).") + if num_topk < 1: + raise ValueError(f"ep_bootstrap requires num_topk >= 1 (got {num_topk}).") + if zero_copy and eager: + raise ValueError("ep_bootstrap: zero_copy requires recv_capacity_per_rank") + if drop_on_overflow and eager: + raise ValueError("ep_bootstrap: drop_on_overflow requires recv_capacity_per_rank") _check_nccl_runtime_version() if zero_copy: warnings.warn( @@ -127,49 +160,78 @@ def ep_bootstrap( str(ep_group.group_name), int(num_experts), int(max_tokens_per_rank), - int(recv_capacity_per_rank), + # Eager mode (recv_capacity_per_rank=None) sizes recv buffers per routing, + # so the group uses the library-derived bound (0 = NCCL_EP_AUTO). + int(recv_capacity_per_rank or 0), int(hidden_dim), int(max_num_sms), max_token_dtype, bool(zero_copy), + int(num_topk), + bool(drop_on_overflow), ) _BOOTSTRAPPED = True _EP_GROUP = ep_group + _EAGER = bool(eager) if not _ATEXIT_REGISTERED: atexit.register(_atexit_finalize) _ATEXIT_REGISTERED = True +def is_ep_bootstrapped() -> bool: + """Whether EP has been initialized in this process.""" + return _BOOTSTRAPPED + + def ep_finalize() -> None: """Optional explicit EP teardown; idempotent. An atexit handler covers normal interpreter shutdown, so most users do not need to call this. Call it explicitly only before ``dist.destroy_process_group()``, since the borrowed NCCL comm becomes - invalid once the PG is destroyed. + invalid once the PG is destroyed. This also releases the symm-mem pool, so + a caller that used ``symm_mem_alloc(use_pool=True)`` does not need a separate + ``release_symm_mem_pool()`` before destroying the PG. """ - global _BOOTSTRAPPED, _EP_GROUP + global _BOOTSTRAPPED, _EP_GROUP, _EAGER if not _BOOTSTRAPPED: return try: + # Deregister pooled symm-mem windows while the group's comm is still valid. + release_symm_mem_pool() tex.ep_finalize() finally: _BOOTSTRAPPED = False _EP_GROUP = None + _EAGER = False + + +def is_symm_backed(t: torch.Tensor) -> bool: + """Whether ``t`` is symm-mem-backed on the EP group. Prefer torch's local ``is_symm_mem_tensor`` + when the build provides it (no collective, no exception); otherwise fall back to the rendezvous + probe the C++ ep kernel uses (``maybe_make_window``): cached for an already-registered tensor, + raises for a plain one.""" + from torch.distributed import _symmetric_memory as _symm + + if hasattr(_symm, "is_symm_mem_tensor"): + return bool(_symm.is_symm_mem_tensor(t)) + if _EP_GROUP is None: + raise RuntimeError( + "is_symm_backed called before ensure_nccl_ep_bootstrapped(); no EP group registered." + ) + try: + _symm.rendezvous(t, _EP_GROUP.group_name) + return True + except Exception: # pylint: disable=broad-exception-caught + return False # Buffer class EpBuffer: - """Per-microbatch EP layer state holding handle_mem and token_counts. + """Per-microbatch EP layer state: handle_mem, tokens_per_expert, and shape/dtype config. Use one EpBuffer per concurrently-in-flight call (e.g. per PP-1F1B microbatch). - - In zero-copy mode the buffer owns the symm-mem buffers the one-sided path - requires: the dispatch recv outputs (recv_tokens, recv_topk_weights) and the - combine backward grad target. One set per buffer, so each layer/microbatch is - isolated. In normal mode these are None and allocated in-flight instead (recv - outputs in the dispatch forward, the combine grad in the backward). """ __slots__ = ( @@ -182,82 +244,75 @@ class EpBuffer: "num_local_experts", "payload_dtype", "device", - "token_counts", + "tokens_per_expert", "zero_copy", - "recv_tokens_symm_buf", - "recv_topk_weights_symm_buf", - "grad_expert_out_symm_buf", + "eager", + "total_recv_tokens", + "_host_total_recv_tokens", + "dispatch_fwd_quant_recipe", + "combine_bwd_quant_recipe", ) - def _alloc_symm_buffers(self) -> None: - """Fill in buffer-owned symm-mem buffers the caller did not supply. - recv_topk_weights is always owned. In normal mode caller-supplied - tensors are kept as-is and the rest stay None (allocated in-flight).""" - if not self.zero_copy: - self.recv_topk_weights_symm_buf = None - return - if _EP_GROUP is None: - raise RuntimeError( - "ep_bootstrap must be called before constructing a zero-copy EpBuffer" - ) - rc, h = self.recv_capacity_per_rank, self.hidden_dim - # Persistent across microbatches; keep resident under CPU offloading. - self.recv_topk_weights_symm_buf = symm_mem_alloc( - (rc,), torch.float32, _EP_GROUP, device=self.device - ) - mark_not_offload(self.recv_topk_weights_symm_buf) - if self.recv_tokens_symm_buf is None: - self.recv_tokens_symm_buf = symm_mem_alloc( - (rc, h), self.payload_dtype, _EP_GROUP, device=self.device - ) - mark_not_offload(self.recv_tokens_symm_buf) - if self.grad_expert_out_symm_buf is None: - self.grad_expert_out_symm_buf = symm_mem_alloc( - (rc, h), self.payload_dtype, _EP_GROUP, device=self.device - ) - mark_not_offload(self.grad_expert_out_symm_buf) - def __init__( self, top_k: int, max_tokens_per_rank: int, - recv_capacity_per_rank: int, hidden_dim: int, num_local_experts: int, + recv_capacity_per_rank: Optional[int] = None, alignment: int = 0, payload_dtype: torch.dtype = torch.bfloat16, device: Optional[torch.device] = None, - dispatch_recv_tokens: Optional[torch.Tensor] = None, - combine_grad_expert_out: Optional[torch.Tensor] = None, + dispatch_fwd_quant_recipe: Optional["Recipe"] = None, + combine_bwd_quant_recipe: Optional["Recipe"] = None, ) -> None: - """Pass ``dispatch_recv_tokens`` (dispatch recv output) and/or - ``combine_grad_expert_out`` (combine backward grad target) to use caller-owned - buffers; the buffer then skips allocating them. Both must be symm-mem-backed - under zero-copy. Whatever is left None is buffer-owned (zero-copy) or allocated - in-flight (normal mode). recv_topk_weights is always owned by the buffer.""" + if not _BOOTSTRAPPED: + raise RuntimeError("EpBuffer requires ep_bootstrap() to be called first.") if device is None: device = torch.device("cuda", torch.cuda.current_device()) alignment = int(alignment) if alignment > 1 and (alignment & (alignment - 1)) != 0: raise ValueError(f"alignment must be 0, 1, or a power of two (got {alignment}).") + self.eager = _EAGER + if not self.eager and recv_capacity_per_rank is None: + raise ValueError( + "EpBuffer requires recv_capacity_per_rank unless the EP group was " + "bootstrapped in eager mode (recv_capacity_per_rank omitted)." + ) self.top_k = int(top_k) self.alignment = alignment self.max_tokens_per_rank = int(max_tokens_per_rank) - self.recv_capacity_per_rank = int(recv_capacity_per_rank) + self.recv_capacity_per_rank = ( + None if recv_capacity_per_rank is None else int(recv_capacity_per_rank) + ) self.hidden_dim = int(hidden_dim) self.num_local_experts = int(num_local_experts) self.payload_dtype = payload_dtype self.device = device self.zero_copy = bool(tex.ep_get_zero_copy()) - self.recv_tokens_symm_buf = dispatch_recv_tokens - self.grad_expert_out_symm_buf = combine_grad_expert_out + self.dispatch_fwd_quant_recipe = dispatch_fwd_quant_recipe + self.combine_bwd_quant_recipe = combine_bwd_quant_recipe size_bytes = tex.ep_handle_mem_size(self.top_k, self.alignment) self.handle_mem = torch.empty(int(size_bytes), dtype=torch.uint8, device=device) - self.token_counts = torch.empty(self.num_local_experts, dtype=torch.int32, device=device) + self.tokens_per_expert = torch.empty( + self.num_local_experts, dtype=torch.int64, device=device + ) # Persistent tensor; keep resident if activation CPU offloading is on. mark_not_offload(self.handle_mem) - self._alloc_symm_buffers() + # Per-step recv-token total (int64 [1]), written by ep_prepare. + # Eager mode uses it to size the recv outputs; graph mode reads it after + # replay to detect overflow past recv_capacity_per_rank. + if self.eager: + # Eager reads this on the host every dispatch. Pinned host memory lets the + # prepare kernel store the total directly to RAM (UVA), so the readback is a + # plain CPU load after one stream sync instead of a pageable D2H round trip. + self.total_recv_tokens = torch.empty(1, dtype=torch.int64, pin_memory=True) + else: + self.total_recv_tokens = torch.empty(1, dtype=torch.int64, device=device) + mark_not_offload(self.total_recv_tokens) + # Host mirror of total_recv_tokens, set by ep_prepare in eager mode. + self._host_total_recv_tokens: Optional[int] = None # torch.library custom ops (so they don't graph-break under torch.compile) @@ -267,17 +322,18 @@ def __init__( @torch.library.custom_op( f"{_LIB}::prepare", - mutates_args=("handle_mem", "token_counts"), + mutates_args=("handle_mem", "tokens_per_expert", "total_recv_tokens"), device_types="cuda", ) def _prepare_op( handle_mem: torch.Tensor, top_k: int, topk_idx: torch.Tensor, - token_counts: torch.Tensor, + tokens_per_expert: torch.Tensor, alignment: int, + total_recv_tokens: torch.Tensor, ) -> None: - tex.ep_prepare(handle_mem, topk_idx, token_counts, top_k, alignment) + tex.ep_prepare(handle_mem, topk_idx, tokens_per_expert, top_k, alignment, total_recv_tokens) @_prepare_op.register_fake @@ -287,7 +343,7 @@ def _(*_args, **_kw): @torch.library.custom_op( f"{_LIB}::dispatch", - mutates_args=("recv_tokens", "recv_topk_weights"), + mutates_args=("recv_tokens", "recv_topk_weights", "recv_scale_inv"), device_types="cuda", ) def _dispatch_op( @@ -297,8 +353,19 @@ def _dispatch_op( topk_weights: torch.Tensor, recv_tokens: torch.Tensor, recv_topk_weights: torch.Tensor, + tokens_scale_inv: Optional[torch.Tensor] = None, + recv_scale_inv: Optional[torch.Tensor] = None, ) -> None: - tex.ep_dispatch(handle_mem, topk_idx, tokens, topk_weights, recv_tokens, recv_topk_weights) + tex.ep_dispatch( + handle_mem, + topk_idx, + tokens, + topk_weights, + recv_tokens, + recv_topk_weights, + tokens_scale_inv, + recv_scale_inv, + ) @_dispatch_op.register_fake @@ -346,15 +413,17 @@ def _(*_args, **_kw): @torch.library.custom_op( f"{_LIB}::combine_bwd", - mutates_args=("grad_expert_out",), + mutates_args=("grad_expert_out", "grad_expert_out_scale_inv"), device_types="cuda", ) def _combine_bwd_op( handle_mem: torch.Tensor, grad: torch.Tensor, grad_expert_out: torch.Tensor, + grad_scale_inv: Optional[torch.Tensor] = None, + grad_expert_out_scale_inv: Optional[torch.Tensor] = None, ) -> None: - tex.ep_combine_bwd(handle_mem, grad, grad_expert_out) + tex.ep_combine_bwd(handle_mem, grad, grad_expert_out, grad_scale_inv, grad_expert_out_scale_inv) @_combine_bwd_op.register_fake @@ -367,13 +436,29 @@ def _(*_args, **_kw): def ep_prepare(buffer: "EpBuffer", topk_idx: torch.Tensor) -> torch.Tensor: """AllGather the routing map; fills ``buffer.handle_mem`` and returns - ``buffer.token_counts`` (int32, shape [num_local_experts]). topk_idx must + ``buffer.tokens_per_expert`` (int64, shape [num_local_experts]). topk_idx must be int32 or int64. + + Also fills ``buffer.total_recv_tokens`` (int64 [1]; pinned host memory in eager + mode, device tensor otherwise) with the per-step recv total; eager mode reads it + on the host to size the recv outputs, graph mode reads it device-side to detect + overflow. """ torch.ops.transformer_engine_ep.prepare( - buffer.handle_mem, buffer.top_k, topk_idx, buffer.token_counts, buffer.alignment + buffer.handle_mem, + buffer.top_k, + topk_idx, + buffer.tokens_per_expert, + buffer.alignment, + buffer.total_recv_tokens, ) - return buffer.token_counts + if buffer.eager: + # total_recv_tokens is pinned host memory stored by the prepare kernel; a CPU + # tensor's .item() does not synchronize, so sync the stream first, then the + # read is a free CPU load (no D2H copy). + torch.cuda.current_stream().synchronize() + buffer._host_total_recv_tokens = int(buffer.total_recv_tokens.item()) + return buffer.tokens_per_expert def _ep_dispatch_raw( @@ -399,56 +484,102 @@ def _ep_combine_raw(buffer: "EpBuffer", expert_out: torch.Tensor, result: torch. class _EpDispatch(torch.autograd.Function): - """Autograd prepare+dispatch; bwd uses user-supplied grad inputs as-is.""" + """Autograd dispatch; caller runs prepare first. bwd uses user-supplied grad inputs as-is.""" @staticmethod def forward( # type: ignore[override] ctx, handle_mem: torch.Tensor, - top_k: int, - alignment: int, - recv_tokens: torch.Tensor, - recv_topk_weights: torch.Tensor, - token_counts: torch.Tensor, + recv_tokens: Optional[torch.Tensor], + recv_topk_weights: Optional[torch.Tensor], topk_idx: torch.Tensor, tokens: torch.Tensor, topk_weights: torch.Tensor, + tokens_scale_inv: Optional[torch.Tensor] = None, + token_counts: Optional[torch.Tensor] = None, + num_recv_tokens: Optional[int] = None, + payload_dtype: torch.dtype = torch.bfloat16, ): - """Prepare + dispatch fwd.""" - torch.ops.transformer_engine_ep.prepare( - handle_mem, top_k, topk_idx, token_counts, alignment - ) + """Dispatch fwd; prepare must have run into ``handle_mem`` beforehand. When scales are set + (MXFP8 for now), ``tokens`` is the quantized tensor kept as the autograd operand so grad + reaches the pre-quant input. Recv outputs are carved/allocated here: a caller may supply + ``recv_tokens`` / ``recv_topk_weights``, else they are sized to ``num_recv_tokens``.""" + is_scaled = tokens_scale_inv is not None + tokens_data = tokens._rowwise_data if isinstance(tokens, QuantizedTensor) else tokens + assert tokens_data.dim() == 2, "EP dispatch tokens must be 2D [num_tokens, hidden]" + hidden = tokens_data.shape[-1] + device = tokens_data.device + zero_copy = tex.ep_get_zero_copy() + + recv_scale_inv = None + if is_scaled: + if tokens._fp8_dtype != tex.DType.kFloat8E4M3: + raise NotImplementedError("EP dispatch supports only E4M3 MXFP8 tokens for now.") + # recv data + scales share one buffer (data then scales); carve or allocate it here. + recv_tokens, recv_scale_inv = _scale_alloc_io( + recv_tokens, + num_recv_tokens, + hidden, + tokens_scale_inv.shape[-1], + tokens_data.dtype, + tokens_scale_inv.dtype, + device, + zero_copy, + ) + # Reinterpret byte-backed FP8 data as the fp8 dtype so the backend sees a scaled tensor. + dispatch_tokens = tokens_data.view(torch.float8_e4m3fn) + dispatch_recv = recv_tokens.view(torch.float8_e4m3fn) + else: + if recv_tokens is None: + recv_tokens = _alloc_io((num_recv_tokens, hidden), payload_dtype, device, zero_copy) + dispatch_tokens = tokens_data + dispatch_recv = recv_tokens + if recv_topk_weights is None: + recv_topk_weights = _alloc_io((num_recv_tokens,), torch.float32, device, zero_copy) torch.ops.transformer_engine_ep.dispatch( handle_mem, topk_idx, - tokens, + dispatch_tokens, topk_weights, - recv_tokens, + dispatch_recv, recv_topk_weights, + tokens_scale_inv, + recv_scale_inv, ) ctx.save_for_backward(handle_mem) ctx.tokens_shape = tokens.shape - ctx.tokens_dtype = tokens.dtype ctx.topk_weights_shape = topk_weights.shape - ctx.tokens_T_flat = tokens.numel() // tokens.shape[-1] + ctx.num_tokens = tokens_data.shape[0] ctx.topk_T_flat = topk_weights.numel() // topk_weights.shape[-1] ctx.top_k = topk_weights.shape[-1] - ctx.recv_capacity = recv_tokens.shape[0] - ctx.hidden_dim = tokens.shape[-1] - ctx.mark_non_differentiable(token_counts) + ctx.hidden_dim = hidden # Detach so the long-lived buffers aren't tracked as differentiable outputs; - # autograd re-attaches grad_fn pointing back at this Function. - return recv_tokens.detach(), recv_topk_weights.detach(), token_counts + # autograd re-attaches grad_fn pointing back at this Function. For scaled inputs + # the expert-major recv data + scales are wrapped into a per-expert GroupedTensor + # so downstream grouped GEMM and autograd see a proper quantized grouped tensor. + if is_scaled: + recv_out = _make_grouped_mxfp8( + recv_tokens.view(tokens._rowwise_data.dtype), + recv_scale_inv, + token_counts, + tokens._fp8_dtype, + tokens.dtype, + ) + else: + recv_out = recv_tokens.detach() + return recv_out, recv_topk_weights.detach() @staticmethod - def backward(ctx, g_recv_tokens, g_recv_topk_weights, _g_token_counts): # type: ignore[override] + def backward(ctx, g_recv_tokens, g_recv_topk_weights): # type: ignore[override] """Dispatch bwd; normalizes grad-input layout, otherwise passes through.""" (handle_mem,) = ctx.saved_tensors device = handle_mem.device g_recv_tokens = g_recv_tokens.contiguous() g_recv_topk_weights = g_recv_topk_weights.contiguous() + # Dispatch grad follows the recv grad's (high-precision) dtype; the quantizer's STE + # owns the fp8 boundary for scaled inputs. grad_tokens = torch.empty( - ctx.tokens_T_flat, ctx.hidden_dim, dtype=ctx.tokens_dtype, device=device + ctx.num_tokens, ctx.hidden_dim, dtype=g_recv_tokens.dtype, device=device ) grad_topk_weights = torch.empty( ctx.topk_T_flat, ctx.top_k, dtype=torch.float32, device=device @@ -462,28 +593,26 @@ def backward(ctx, g_recv_tokens, g_recv_topk_weights, _g_token_counts): # type: ) return ( None, # handle_mem - None, # top_k - None, # alignment None, # recv_tokens None, # recv_topk_weights - None, # token_counts None, # topk_idx grad_tokens.view(ctx.tokens_shape), grad_topk_weights.view(ctx.topk_weights_shape), + None, # tokens_scale_inv (scales; non-differentiable) + None, # token_counts (per-expert counts; non-differentiable) + None, # num_recv_tokens (sizing scalar) + None, # payload_dtype (sizing scalar) ) class _EpCombine(torch.autograd.Function): - """Autograd combine. + """Autograd combine; bwd scatters the expert_out grad into ``grad_out``. When the caller + supplies it that buffer is used as-is; otherwise it is allocated in the backward from the + symm-mem pool in zero-copy mode (one-sided target) or a plain tensor in normal mode (keeps + allocation torch.compile / CUDA-graph safe and lets autograd own the grad's lifetime). - bwd scatters the expert_out grad into ``grad_symm_buf`` (EpBuffer-owned - symm-mem, one-sided) in zero-copy mode, or into a plain tensor allocated - in-flight here otherwise. The latter keeps allocation torch.compile / - CUDA-graph safe and lets autograd own the grad's lifetime. - - ``grad_symm_buf`` is the backward's scatter target (an output it writes, never - reads), so it is stashed as a plain ctx attribute rather than via - save_for_backward, which would version-track a tensor we mutate. + ``grad_out`` is a write-only scatter target, so it is stashed as a plain ctx attribute rather + than via save_for_backward, which would version-track a tensor we mutate. """ @staticmethod @@ -492,51 +621,205 @@ def forward( # type: ignore[override] handle_mem: torch.Tensor, num_local_tokens: int, hidden_dim: int, - grad_symm_buf: Optional[torch.Tensor], + grad_out: Optional[torch.Tensor], expert_out: torch.Tensor, + bwd_quant_recipe=None, + token_counts: Optional[torch.Tensor] = None, ): - """Combine fwd; stashes the bwd grad target or expert_out shape to size it.""" + """Combine fwd; stashes the bwd grad target or expert_out shape to size it. When + ``bwd_quant_recipe`` is set, the backward sends the result-grad as MXFP8.""" device = expert_out.device result = torch.empty(num_local_tokens, hidden_dim, dtype=expert_out.dtype, device=device) torch.ops.transformer_engine_ep.combine(handle_mem, expert_out, result) ctx.save_for_backward(handle_mem) - ctx.grad_symm_buf = grad_symm_buf - if grad_symm_buf is None: - ctx.expert_out_shape = expert_out.shape - ctx.expert_out_dtype = expert_out.dtype - ctx.device = device + ctx.grad_out = grad_out + ctx.bwd_quant_recipe = bwd_quant_recipe + ctx.token_counts = token_counts + ctx.expert_out_shape = expert_out.shape + ctx.expert_out_dtype = expert_out.dtype + ctx.device = device return result @staticmethod def backward(ctx, g_result): # type: ignore[override] - """Combine bwd; scatters the result grad into the grad target.""" + """Combine bwd; scatters the result-grad to expert positions. High-precision sends the grad + as-is; a quantized recipe (MXFP8 today) quantizes it and returns the expert_out grad as a + per-expert GroupedTensor.""" if not g_result.is_contiguous(): g_result = g_result.contiguous() (handle_mem,) = ctx.saved_tensors - grad_expert_out = ctx.grad_symm_buf - if grad_expert_out is None: - grad_expert_out = torch.empty( - ctx.expert_out_shape, dtype=ctx.expert_out_dtype, device=ctx.device + + if ctx.bwd_quant_recipe is None: + grad_expert_out = ctx.grad_out + if grad_expert_out is None: + grad_expert_out = _alloc_io( + ctx.expert_out_shape, ctx.expert_out_dtype, ctx.device, tex.ep_get_zero_copy() + ) + torch.ops.transformer_engine_ep.combine_bwd(handle_mem, g_result, grad_expert_out) + else: + mx, g_scale_inv = _quantize_mxfp8(g_result) + g_data = mx._rowwise_data + recv_pr, hidden = ctx.expert_out_shape[0], ctx.expert_out_shape[-1] + ge_data, ge_scale_inv = _scale_alloc_io( + ctx.grad_out, + recv_pr, + hidden, + g_scale_inv.shape[-1], + g_data.dtype, + g_scale_inv.dtype, + ctx.device, + tex.ep_get_zero_copy(), + ) + # The backend keys on the fp8 scaling mode; reinterpret the byte-backed data as fp8. + torch.ops.transformer_engine_ep.combine_bwd( + handle_mem, + g_data.view(torch.float8_e4m3fn), + ge_data.view(torch.float8_e4m3fn), + g_scale_inv, + ge_scale_inv, ) - torch.ops.transformer_engine_ep.combine_bwd(handle_mem, g_result, grad_expert_out) + grad_expert_out = _make_grouped_mxfp8( + ge_data, ge_scale_inv, ctx.token_counts, mx._fp8_dtype, ctx.expert_out_dtype + ) + return ( None, # handle_mem None, # num_local_tokens None, # hidden_dim - None, # grad_symm_buf + None, # grad_out grad_expert_out, + None, # bwd_quant_recipe + None, # token_counts ) # Public high-level wrappers -# NCCL EP currently only supports bfloat16 payload tensors. +# NCCL EP inputs are bfloat16; MXFP8 is applied internally via the buffer's dispatch_fwd_quant_recipe. def _require_bf16(name: str, t: torch.Tensor) -> None: if t.dtype is not torch.bfloat16: raise NotImplementedError( - f"NCCL EP currently supports only bfloat16 payloads; got {name}.dtype={t.dtype}." + "NCCL EP currently supports only bfloat16 or MXFP8 payloads; got" + f" {name}.dtype={t.dtype}." + ) + + +def _alloc_io(shape, dtype: torch.dtype, device, zero_copy: bool) -> torch.Tensor: + """Allocate a dispatch/combine IO tensor the caller did not supply: from the symm-mem pool in + zero-copy mode (auto-registered segment, lifecycle managed by torch refcount), else plain. + + The zero-copy pool path is not CUDA-graph capturable; supply persistent recv_tokens / grad_out + buffers to capture a graph.""" + if zero_copy: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "EP zero-copy pool allocation is not CUDA-graph capturable; supply persistent " + "recv_tokens / grad_out buffers (allocated once via symm_mem_alloc) before capture." + ) + t = symm_mem_alloc(shape, dtype, _EP_GROUP, device=device, use_pool=True) + # symm-mem storage is non-resizable; exempt it from CPU activation offloading (which + # releases via storage.resize_(0)). Matters for bf16 recv_tokens (the saved activation). + mark_not_offload(t) + return t + return torch.empty(*shape, dtype=dtype, device=device) + + +def _quantize_mxfp8(x: torch.Tensor): + """Quantize a high-precision tensor to MXFP8 and return ``(quantized_tensor, scale_inv)`` where + ``scale_inv`` is the compact ``[T, H/block]`` scale the EP backend routes. The quantized tensor + is returned so callers can keep it as the autograd operand; its ``_rowwise_data`` is the fp8 + payload and ``scale_inv.shape[-1]`` the scale-column count. EP routes and returns E4M3 data in + both directions, so quantize to E4M3 regardless of pass. Strips the GEMM scale row padding to + the compact ``[T, H/block]`` layout; requires a 16-byte-aligned scale row.""" + from .constants import MXFP8_BLOCK_SCALING_SIZE + from .tensor.mxfp8_tensor import MXFP8Quantizer + + mx = MXFP8Quantizer(tex.DType.kFloat8E4M3, rowwise=True, columnwise=False).quantize(x) + if mx._with_gemm_swizzled_scales: + raise RuntimeError( + "internal MXFP8 quantization produced swizzled scales; EP dispatch needs compact." + ) + data = mx._rowwise_data + scale_inv = mx._rowwise_scale_inv + if data is None or scale_inv is None: + raise ValueError("MXFP8 tokens must carry rowwise data and scale_inv for EP dispatch.") + t_flat = x.shape[0] + hidden = x.shape[-1] + cols = hidden // MXFP8_BLOCK_SCALING_SIZE + # The backend forwards each token's scale row with a 16-byte-aligned store, so the row + # (cols * dtype bytes) must be a multiple of 16. + scale_row_bytes = cols * scale_inv.element_size() + if scale_row_bytes % 16 != 0: + raise ValueError( + f"MXFP8 dispatch requires a 16-byte-aligned scale row; hidden={hidden} gives " + f"{scale_row_bytes} bytes. Use a hidden size that is a multiple of " + f"{16 * MXFP8_BLOCK_SCALING_SIZE}." ) + # scale_inv is 2D [round_up(T, 128), cols]; drop the row padding to the logical [T, H/block] + # the backend expects. cols is a multiple of 4 (16-byte row), so no column padding and the + # slice stays contiguous; assert rather than force a copy. + scale_inv = scale_inv[:t_flat, :cols] + if not scale_inv.is_contiguous(): + raise ValueError( + "MXFP8 dispatch requires compact contiguous scales [T, H/block]; got a " + f"non-contiguous [{t_flat}, {cols}] slice." + ) + return mx, scale_inv + + +def _scale_alloc_io(buf, rows, data_cols, scale_cols, data_dtype, scale_dtype, device, zero_copy): + """Block-scaled output data + scale buffers, each ``rows`` tall, laid out back-to-back + (``[rows, data_cols]`` data of ``data_dtype`` then ``[rows, scale_cols]`` scales of + ``scale_dtype``). Carve both from a single caller ``buf`` when it is large enough, so one + symm-mem window backs both views; else allocate them (symm-mem pool under zero-copy, else + plain). Recipe-agnostic: byte sizes come from the element sizes.""" + data_bytes = rows * data_cols * torch.empty((), dtype=data_dtype).element_size() + scale_bytes = rows * scale_cols * torch.empty((), dtype=scale_dtype).element_size() + if buf is not None: + # Reinterpret in place; a non-contiguous buf would force a copy and leave the caller's + # buffer unwritten, so require contiguous and view rather than reshape. + if not buf.is_contiguous(): + raise ValueError("scaled output buffer must be contiguous.") + flat = buf.view(-1).view(torch.uint8) + if flat.numel() < data_bytes + scale_bytes: + raise ValueError( + f"scaled output buffer too small: need {data_bytes + scale_bytes} bytes " + f"(data + scales), got {flat.numel()}." + ) + data = flat[:data_bytes].view(data_dtype).reshape(rows, data_cols) + scale_inv = ( + flat[data_bytes : data_bytes + scale_bytes].view(scale_dtype).reshape(rows, scale_cols) + ) + return data, scale_inv + data = _alloc_io((rows, data_cols), data_dtype, device, zero_copy) + scale_inv = _alloc_io((rows, scale_cols), scale_dtype, device, zero_copy) + return data, scale_inv + + +def _make_grouped_mxfp8(data, scale_inv, token_counts, fp8_dtype, fake_dtype): + """Wrap expert-major MXFP8 recv data + compact e8m0 scales as a per-expert ``GroupedTensor``. + + ``token_counts`` (int64 [num_local_experts]) is the padded per-expert row counts (128-aligned), + used as the group sizes. Grouping is device-side (first_dims/tensor_offsets), so the counts never + sync to host; the outer shape is the static recv capacity, bounded per expert by first_dims. + """ + from .tensor.grouped_tensor import GroupedTensor + from .tensor.mxfp8_tensor import MXFP8Quantizer + + assert data.dim() == 2, "recv data must be 2D [capacity_rows, hidden]" + capacity_rows, hidden = data.shape + quantizer = MXFP8Quantizer(fp8_dtype, rowwise=True, columnwise=False) + return GroupedTensor( + shape=(capacity_rows, hidden), + dtype=fake_dtype, + num_tensors=token_counts.numel(), + quantizer=quantizer, + data=data.reshape(-1).detach(), + scale_inv=scale_inv.reshape(-1).detach(), + first_dims=token_counts, + tensor_offsets=tex.splits_to_offsets(token_counts, hidden), + ) def ep_dispatch( @@ -544,44 +827,76 @@ def ep_dispatch( tokens: torch.Tensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor, + *, + recv_tokens: Optional[torch.Tensor] = None, + recv_topk_weights: Optional[torch.Tensor] = None, ): - """Prepare + dispatch with autograd. topk_idx must be int32 or int64. + """Prepare + dispatch with autograd. ``tokens`` is bfloat16; ``topk_idx`` is int32 or int64. + + When the buffer's ``dispatch_fwd_quant_recipe`` is set (``MXFP8BlockScaling`` only for now), tokens + are quantized internally and recv is returned as a per-expert ``GroupedTensor``; otherwise recv + stays bfloat16. A pre-quantized ``tokens`` is not accepted. - recv_tokens comes from the EpBuffer (caller-supplied or buffer-owned under - zero-copy) or is allocated in-flight (normal mode). recv_topk_weights is always - owned by the buffer. Returns (recv_tokens, recv_topk_weights, token_counts); - token_counts is non-diff. + ``recv_tokens`` / ``recv_topk_weights`` are the recv outputs: pass caller-owned buffers + (symm-mem-backed under zero-copy) or leave them None to allocate. For MXFP8 the recv data and + scales share ``recv_tokens`` (data then scales), so size it to at least + ``recv_capacity_per_rank * (hidden + hidden/block)`` bytes. Eager mode sizes the recv outputs + per step and forbids caller-supplied buffers. Under zero-copy, leaving them None allocates from + the symm-mem pool, which is not CUDA-graph capturable; pass persistent buffers to capture a graph. + + Returns (recv_tokens, recv_topk_weights, tokens_per_expert); tokens_per_expert is non-diff. See + ``buffer.total_recv_tokens`` for the per-step recv total. """ - _require_bf16("tokens", tokens) if topk_weights.dtype is not torch.float32: raise TypeError( f"topk_weights must be float32; got dtype={topk_weights.dtype}. " "Cast with topk_weights.float() before calling." ) - recv_tokens = buffer.recv_tokens_symm_buf - if recv_tokens is None: - recv_tokens = torch.empty( - buffer.recv_capacity_per_rank, - buffer.hidden_dim, - dtype=buffer.payload_dtype, - device=buffer.device, + if isinstance(tokens, QuantizedTensor): + raise NotImplementedError( + "NCCL EP dispatch takes a bfloat16 input and quantizes internally when the buffer's " + "dispatch_fwd_quant_recipe is set; a pre-quantized tensor is not accepted." + ) + _require_bf16("tokens", tokens) + if buffer.eager and (recv_tokens is not None or recv_topk_weights is not None): + raise ValueError( + "eager mode sizes the recv outputs from the per-step recv-token total " + "and cannot use caller-supplied recv_tokens / recv_topk_weights" ) - recv_topk_weights = ( - buffer.recv_topk_weights_symm_buf - if buffer.zero_copy - else torch.empty(buffer.recv_capacity_per_rank, dtype=torch.float32, device=buffer.device) + + # Prepare (routing AllGather) up front so the recv outputs can be sized; in + # eager mode ep_prepare also host-syncs this step's recv-token total. + tokens_per_expert = ep_prepare(buffer, topk_idx) + num_recv_tokens = ( + buffer._host_total_recv_tokens if buffer.eager else buffer.recv_capacity_per_rank ) - return _EpDispatch.apply( + + tokens_scale_inv = None + if buffer.dispatch_fwd_quant_recipe is not None: + from ..common.recipe import MXFP8BlockScaling + + if not isinstance(buffer.dispatch_fwd_quant_recipe, MXFP8BlockScaling): + raise NotImplementedError( + "EP block-scaled dispatch supports MXFP8BlockScaling only; got " + f"{type(buffer.dispatch_fwd_quant_recipe).__name__}." + ) + # Quantize here (not in forward) so the quantized tensor stays the autograd operand and grad + # reaches the pre-quant input; forward then carves the recv buffers and routes. + tokens, tokens_scale_inv = _quantize_mxfp8(tokens) + + recv_tokens, recv_topk_weights = _EpDispatch.apply( buffer.handle_mem, - buffer.top_k, - buffer.alignment, recv_tokens, recv_topk_weights, - buffer.token_counts, topk_idx, tokens, topk_weights, + tokens_scale_inv, + tokens_per_expert, + num_recv_tokens, + buffer.payload_dtype, ) + return recv_tokens, recv_topk_weights, tokens_per_expert def ep_combine( @@ -589,22 +904,45 @@ def ep_combine( expert_out: torch.Tensor, *, num_local_tokens: Optional[int] = None, + grad_out: Optional[torch.Tensor] = None, ): """Combine with autograd; caller pre-applies topk weighting. - The backward scatters the expert_out grad into the EpBuffer grad target - (caller-supplied or buffer-owned under zero-copy), or a tensor allocated - in-flight (normal mode). Result shape is (num_local_tokens, buffer.hidden_dim); - defaults to buffer.max_tokens_per_rank rows. + ``expert_out`` is the combine input (symm-mem-backed under zero-copy). ``grad_out`` is the + backward's grad target: pass a caller-owned buffer or leave it None to allocate. For MXFP8 the + grad data and scales share ``grad_out`` (data then scales), so size it to at least + ``recv_capacity_per_rank * (hidden + hidden/block)`` bytes (non-zero-copy only). Eager mode sizes + the grad target per step and forbids a caller-supplied buffer. Under zero-copy, leaving it None + allocates from the symm-mem pool, which is not CUDA-graph capturable; pass a persistent buffer to + capture a graph. Result shape is (num_local_tokens, hidden_dim); num_local_tokens defaults to + buffer.max_tokens_per_rank. """ _require_bf16("expert_out", expert_out) + if buffer.eager and grad_out is not None: + raise ValueError( + "eager mode sizes the combine grad target per step and cannot use a " + "caller-supplied grad_out" + ) if num_local_tokens is None: num_local_tokens = buffer.max_tokens_per_rank - grad_expert_out = buffer.grad_expert_out_symm_buf + # When combine_bwd_quant_recipe is set the combine backward sends the result-grad over the + # wire as MXFP8 and returns the expert_out grad as a GroupedTensor. + bwd_quant_recipe = None + if buffer.combine_bwd_quant_recipe is not None: + from ..common.recipe import MXFP8BlockScaling + + if not isinstance(buffer.combine_bwd_quant_recipe, MXFP8BlockScaling): + raise NotImplementedError( + "EP combine backward supports MXFP8BlockScaling only; got " + f"{type(buffer.combine_bwd_quant_recipe).__name__}." + ) + bwd_quant_recipe = buffer.combine_bwd_quant_recipe return _EpCombine.apply( buffer.handle_mem, num_local_tokens, buffer.hidden_dim, - grad_expert_out, + grad_out, expert_out, + bwd_quant_recipe, + buffer.tokens_per_expert, ) diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 2e966c7e94..bcbf9c5eb2 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -38,6 +38,12 @@ _T = TypeVar("_T") SingleOrTuple = Union[_T, Tuple[_T, ...]] +_CAPTURE_TIME_HOOK_NAMES = ( + "forward_pre_hooks", + "forward_hooks", + "backward_pre_hooks", + "backward_hooks", +) def set_capture_start() -> None: @@ -98,9 +104,50 @@ def _graph_context_wrapper(*args, **kwargs): gc.enable() +def _canonicalize_capture_time_hooks( + num_callables: int, + capture_time_hooks: Optional[List[Optional[Dict[str, Dict]]]], +) -> List[Dict[str, Dict]]: + """Fill defaults in capture_time_hooks.""" + if capture_time_hooks is None: + capture_time_hooks = [None] * num_callables + if len(capture_time_hooks) != num_callables: + raise ValueError( + f"capture_time_hooks has {len(capture_time_hooks)} entries, " + f"but there are {num_callables} callables." + ) + + canonicalized = [] + for callable_idx in range(num_callables): + hooks = capture_time_hooks[callable_idx] + if hooks is None: + hooks = {} + unexpected_keys = set(hooks.keys()) - set(_CAPTURE_TIME_HOOK_NAMES) + if unexpected_keys: + raise ValueError(f"Found unexpected keys in capture_time_hooks ({unexpected_keys}).") + canonicalized.append( + {hook_name: hooks.get(hook_name, {}) for hook_name in _CAPTURE_TIME_HOOK_NAMES} + ) + + return canonicalized + + +def _run_capture_time_hooks( + capture_time_hooks: List[Dict[str, Dict]], + callable_idx: int, + hook_name: str, + module: torch.nn.Module, +) -> None: + """Run non-capturable hooks outside CUDA graph capture.""" + for hook in capture_time_hooks[callable_idx][hook_name].values(): + if hook(module) is not None: + raise RuntimeError(f"capture_time_hooks {hook_name} must not return a value.") + + def _make_graphed_callables( callables: SingleOrTuple[Callable], sample_args: SingleOrTuple[Tuple[torch.Tensor, ...]], + *, num_warmup_iters: int = 3, allow_unused_input: bool = False, cache_quantized_params: bool = False, @@ -110,8 +157,10 @@ def _make_graphed_callables( pool: Optional[Tuple[int, ...]] = None, retain_graph_in_backward: bool = False, _reuse_graph_input_output_buffers: bool = False, + clone_param_grads_on_return: bool = True, pre_warmup_hook: Optional[Callable] = None, post_warmup_hook: Optional[Callable] = None, + capture_time_hooks: List[Dict[str, Dict]], ) -> SingleOrTuple[Callable]: """ Helper method for `make_graphed_callables` @@ -333,13 +382,18 @@ def _make_graphed_callables( if isinstance(c, torch.nn.Module): if not ( len(c._backward_hooks) == 0 + and len(c._backward_pre_hooks) == 0 and len(c._forward_hooks) == 0 and len(c._forward_pre_hooks) == 0 ): raise RuntimeError( "Modules must not have hooks registered at the time they are passed. " + "However, registering hooks on modules after passing them " - + "through make_graphed_callables is allowed." + + "through make_graphed_callables is allowed. " + + "If you have to use hooks during capture time, you can provide them " + + "in the capture_time_hooks argument, and they will be executed outside " + + "the CUDA graph capture context, meaning they will not be recorded into " + + "the graph and will not be replayed." ) if not all(b.requires_grad is False for b in c.buffers()): raise RuntimeError( @@ -404,6 +458,24 @@ def _make_graphed_callables( bwd_dw_graphs = [torch.cuda.CUDAGraph() for _ in range(len(flatten_sample_args))] graph_callables = [None for _ in range(len(flatten_sample_args))] + def _returned_param_grad_clone_slots(static_grad_inputs, module_params): + """Snapshot static grad slots that need clones before returning from Graphed.backward.""" + if not clone_param_grads_on_return: + return (False,) * len(static_grad_inputs) + module_param_start = len(static_grad_inputs) - len(module_params) + # `skip_backward_post_hook` marks parameters whose gradient lifetime is + # managed by the delayed-wgrad module hook. With fused accumulation, the + # returned weight grad is only a dummy and does not carry the actual wgrad, + # which is written directly to `main_grad`. Therefore, these slots do not + # need to be cloned. + return tuple( + idx >= module_param_start + and not getattr( + module_params[idx - module_param_start], "skip_backward_post_hook", False + ) + for idx in range(len(static_grad_inputs)) + ) + # For cases with multiple active RNG states, e.g. TP. if graph_safe_rng_available(): for _, state in get_all_rng_states().items(): @@ -450,113 +522,161 @@ def _make_graphed_callables( visited_te_modules = {} need_bwd_dw_graph = {} + def _run_warmup_forward(func_idx, func, callable_idx): + """Run forward for one callable during warmup; returns flattened outputs.""" + args = sample_args[func_idx] + kwargs = sample_kwargs[func_idx] + + def hook_fn(module, inputs, outputs, func_idx=func_idx): # pylint: disable=unused-argument + modules = set() + if isinstance(module, TransformerEngineBaseModule): + modules.add(module) + # If forward is called on a BasicOperation directly the hook will run + elif isinstance(module, BasicOperation): + modules.add(module) + # If forward is called on a te.ops.Sequential it is not called on its constituent ops + elif isinstance(module, Sequential): + if module._module_groups is None: + raise RuntimeError( + "module._module_groups should have been initialized by warmup" + ) + for module_group in module._module_groups: + if isinstance(module_group, OperationFuser): + for basic_op in module_group._basic_ops: + modules.add(basic_op) + if modules: + if func_idx not in visited_te_modules: + visited_te_modules[func_idx] = modules + else: + visited_te_modules[func_idx].update(modules) + + _run_capture_time_hooks(capture_time_hooks, callable_idx, "forward_pre_hooks", func) + + hooks = [] + for module in func.modules(): + hooks.append(module.register_forward_hook(hook_fn)) + outputs = func(*args, **kwargs) + for hook in hooks: + hook.remove() + + _run_capture_time_hooks(capture_time_hooks, callable_idx, "forward_hooks", func) + + outputs, _ = _tree_flatten(outputs) + return outputs + + def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx): + """Run dgrad backward for one callable during warmup.""" + static_input_surface = per_callable_static_input_surfaces[func_idx] + + inputs = tuple(i for i in static_input_surface if i.requires_grad) + outputs_requiring_grad = tuple(o for o in outputs if o is not None and o.requires_grad) + grad_outputs = tuple(torch.empty_like(o) for o in outputs_requiring_grad) + + _run_capture_time_hooks(capture_time_hooks, callable_idx, "backward_pre_hooks", func) + + with _none_grad_context_wrapper(inputs): + torch.autograd.backward(outputs_requiring_grad, grad_tensors=grad_outputs) + grad_inputs = tuple(input.grad for input in inputs) + + _run_capture_time_hooks(capture_time_hooks, callable_idx, "backward_hooks", func) + + # Filter module params that get None grad from grad_inputs and remove them + # from static_input_surface. This is to ensure that the backward hooks + # registered to these params are not wrongly triggered. + num_required_grad_sample_args = sum( + arg.requires_grad for arg in flatten_sample_args[func_idx] + ) + required_grad_input_idx = [] + for i, arg in enumerate(static_input_surface): + if arg.requires_grad: + required_grad_input_idx.append(i) + module_params_with_grad = [] + for grad_inputs_idx, inputs_idx in enumerate(required_grad_input_idx): + if ( + grad_inputs[grad_inputs_idx] is None + and grad_inputs_idx < num_required_grad_sample_args + ): + if not allow_unused_input: + raise RuntimeError( + "The input tensor requires grad, but the grad is None after backward pass." + ) + elif ( + grad_inputs[grad_inputs_idx] is not None + and grad_inputs_idx >= num_required_grad_sample_args + ): + module_params_with_grad.append(static_input_surface[inputs_idx]) + if len(module_params_with_grad) != len(per_callable_module_params[func_idx]): + if warmup_iter != 0: + raise RuntimeError( + "no-grad params should only be used as inputs in the first warmup" + f" iteration, but found in iteration {warmup_iter}" + ) + per_callable_module_params[func_idx] = tuple(module_params_with_grad) + static_input_surface = flatten_sample_args[func_idx] + tuple(module_params_with_grad) + per_callable_static_input_surfaces[func_idx] = static_input_surface + + # Run wgrad. This is essential for some TE modules when they have + # delay_wgrad_compute enabled. + need_backward_dw = False + for module in visited_te_modules.get(func_idx, set()): + if hasattr(module, "need_backward_dw") and module.need_backward_dw(): + need_backward_dw = True + module.backward_dw() + need_bwd_dw_graph[func_idx] = need_backward_dw + # Run warmup and do the above filtering. # ROCm: reuse warmup stream for graph capture (ROCM-25129) stream = torch.cuda.Stream() with torch.cuda.stream(stream): - for func_idx, func in zip(warmup_func_idx, warmup_func): - args = sample_args[func_idx] - kwargs = sample_kwargs[func_idx] - static_input_surface = per_callable_static_input_surfaces[func_idx] - - def hook_fn( - module, inputs, outputs, func_idx=func_idx - ): # pylint: disable=unused-argument - modules = set() - if isinstance(module, TransformerEngineBaseModule): - modules.add(module) - # If forward is called on a BasicOperation directly the hook will run - elif isinstance(module, BasicOperation): - modules.add(module) - # If forward is called on a te.ops.Sequential it is not called on its constituent ops - elif isinstance(module, Sequential): - if module._module_groups is None: - raise RuntimeError( - "module._module_groups should have been initialized by warmup" - ) - for module_group in module._module_groups: - if isinstance(module_group, OperationFuser): - for basic_op in module_group._basic_ops: - modules.add(basic_op) - if modules: - if func_idx not in visited_te_modules: - visited_te_modules[func_idx] = modules - else: - visited_te_modules[func_idx].update(modules) - - if pre_warmup_hook is not None: - pre_warmup_hook() - for warmup_iter in range(num_warmup_iters): - hooks = [] - for module in func.modules(): - hook = module.register_forward_hook(hook_fn) - hooks.append(hook) - outputs, _ = _tree_flatten(func(*args, **kwargs)) - for hook in hooks: - hook.remove() - if is_training: - inputs = tuple(i for i in static_input_surface if i.requires_grad) - with _none_grad_context_wrapper(inputs): - outputs_requiring_grad = tuple( - o for o in outputs if o is not None and o.requires_grad - ) - torch.autograd.backward( - outputs_requiring_grad, - grad_tensors=tuple(torch.empty_like(o) for o in outputs_requiring_grad), - ) - grad_inputs = tuple(input.grad for input in inputs) + if pre_warmup_hook is not None: + pre_warmup_hook() - # Filter module params that get None grad from grad_inputs and remove them - # from static_input_surface. This is to ensure that the backward hooks - # registered to these params are not wrongly triggered. - num_required_grad_sample_args = sum( - arg.requires_grad for arg in flatten_sample_args[func_idx] - ) - required_grad_input_idx = [] - for i, arg in enumerate(static_input_surface): - if arg.requires_grad: - required_grad_input_idx.append(i) - module_params_with_grad = [] - for grad_inputs_idx, inputs_idx in enumerate(required_grad_input_idx): - if ( - grad_inputs[grad_inputs_idx] is None - and grad_inputs_idx < num_required_grad_sample_args - ): - if not allow_unused_input: - raise RuntimeError( - "The input tensor requires grad, but the grad is None after" - " backward pass." + for warmup_iter in range(num_warmup_iters): + if _order is None: + # All forwards in order, then all backwards in reverse order. + warmup_outputs = [] + for func_idx, func in zip(warmup_func_idx, warmup_func): + outputs = _run_warmup_forward(func_idx, func, func_idx) + warmup_outputs.append((func_idx, func, outputs)) + if is_training: + for func_idx, func, outputs in reversed(warmup_outputs): + _run_warmup_backward(func_idx, func, outputs, warmup_iter, func_idx) + else: + # Follow _order exactly, mirroring the capture phase. + per_fwd_outputs = {} # per_callable_fwd_idx -> flattened outputs + fwd_idx = [0] * num_model_chunks + bwd_idx = [0] * num_model_chunks + for c_id in _order: + if c_id > 0: + # Forward pass for chunk c_id. + m_chunk = c_id - 1 + for l_no in range(_num_layers_per_chunk[m_chunk]): + callable_idx = _prefix_num_layers[m_chunk] + l_no + per_callable_fwd_idx = ( + _prefix_num_layers[m_chunk] * num_microbatches + ) + (fwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + l_no) + func = callables[callable_idx] + outputs = _run_warmup_forward(per_callable_fwd_idx, func, callable_idx) + per_fwd_outputs[per_callable_fwd_idx] = outputs + fwd_idx[m_chunk] += 1 + elif ceil(c_id) == c_id: + # Backward pass for chunk -c_id. + if is_training: + m_chunk = -c_id - 1 + for l_no in reversed(range(_num_layers_per_chunk[m_chunk])): + callable_idx = _prefix_num_layers[m_chunk] + l_no + per_callable_bwd_idx = ( + _prefix_num_layers[m_chunk] * num_microbatches + ) + (bwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + l_no) + func = callables[callable_idx] + outputs = per_fwd_outputs[per_callable_bwd_idx] + _run_warmup_backward( + per_callable_bwd_idx, func, outputs, warmup_iter, callable_idx ) - elif ( - grad_inputs[grad_inputs_idx] is not None - and grad_inputs_idx >= num_required_grad_sample_args - ): - module_params_with_grad.append(static_input_surface[inputs_idx]) - if len(module_params_with_grad) != len(per_callable_module_params[func_idx]): - if warmup_iter != 0: - raise RuntimeError( - "no-grad params should only be used as inputs in the first warmup" - f" iteration, but found in iteration {warmup_iter}" - ) - per_callable_module_params[func_idx] = tuple(module_params_with_grad) - static_input_surface = flatten_sample_args[func_idx] + tuple( - module_params_with_grad - ) - per_callable_static_input_surfaces[func_idx] = static_input_surface - - # Run wgrad. This is essential for some TE modules when they have - # delay_wgrad_compute enabled. - need_backward_dw = False - for module in visited_te_modules.get(func_idx, set()): - if hasattr(module, "need_backward_dw") and module.need_backward_dw(): - need_backward_dw = True - module.backward_dw() - need_bwd_dw_graph[func_idx] = need_backward_dw - else: - grad_inputs = None - del outputs, grad_inputs - if post_warmup_hook is not None: - post_warmup_hook() + bwd_idx[m_chunk] += 1 + + if post_warmup_hook is not None: + post_warmup_hook() torch.cuda.synchronize() # All captures here share a mempool. To avoid replays corrupting each other's memory, @@ -568,6 +688,7 @@ def hook_fn( per_callable_output_unflatten_spec = [None] * len(flatten_sample_args) per_callable_static_grad_outputs = [None] * len(flatten_sample_args) per_callable_static_grad_inputs = [None] * len(flatten_sample_args) + per_callable_returned_param_grad_clone_slots = [None] * len(flatten_sample_args) fwd_idx = [0] * num_model_chunks bwd_idx = [0] * num_model_chunks static_grad_outputs_dict = {} @@ -582,15 +703,31 @@ def hook_fn( # Capture forward graph for model chunk c_id, microbatch fwd_idx[c_id-1] m_chunk = c_id - 1 for l_no in range(_num_layers_per_chunk[m_chunk]): - func = callables[_prefix_num_layers[m_chunk] + l_no] + callable_idx = _prefix_num_layers[m_chunk] + l_no + func = callables[callable_idx] per_callable_fwd_idx = (_prefix_num_layers[m_chunk] * num_microbatches) + ( fwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + l_no ) args = sample_args[per_callable_fwd_idx] kwargs = sample_kwargs[per_callable_fwd_idx] fwd_graph = fwd_graphs[per_callable_fwd_idx] + _run_capture_time_hooks( + capture_time_hooks, + callable_idx, + "forward_pre_hooks", + func, + ) + with _graph_context_wrapper(fwd_graph, stream=stream, pool=mempool): outputs = func(*args, **kwargs) + + _run_capture_time_hooks( + capture_time_hooks, + callable_idx, + "forward_hooks", + func, + ) + flatten_outputs, spec = _tree_flatten(outputs) per_callable_static_outputs[per_callable_fwd_idx] = tuple(flatten_outputs) per_callable_output_unflatten_spec[per_callable_fwd_idx] = spec @@ -601,6 +738,7 @@ def hook_fn( m_chunk = -ceil(c_id) - 1 previous_per_callable_bwd_idx = None for l_no in list(reversed(range(_num_layers_per_chunk[m_chunk]))): + callable_idx = _prefix_num_layers[m_chunk] + l_no per_callable_bwd_idx = (_prefix_num_layers[m_chunk] * num_microbatches) + ( bwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + l_no ) @@ -686,7 +824,15 @@ def hook_fn( torch.empty_like(o) if o is not None and o.requires_grad else None for o in static_outputs ) + grad_inputs: Tuple[Optional[torch.Tensor], ...] = () if is_training: + _run_capture_time_hooks( + capture_time_hooks, + callable_idx, + "backward_pre_hooks", + callables[callable_idx], + ) + inputs = tuple(i for i in static_input_surface if i.requires_grad) with _none_grad_context_wrapper(inputs), _graph_context_wrapper( bwd_graph, stream=stream, pool=mempool @@ -700,6 +846,13 @@ def hook_fn( ) grad_inputs = tuple(input.grad for input in inputs) + _run_capture_time_hooks( + capture_time_hooks, + callable_idx, + "backward_hooks", + callables[callable_idx], + ) + # Constructs a tuple suitable for returning from Graphed.backward: # Pads out the actually-needed grads with Nones in gradient slots for inputs # that don't require grad. I couldn't think of a one-liner for this pattern. @@ -715,6 +868,13 @@ def hook_fn( per_callable_static_grad_outputs[per_callable_bwd_idx] = static_grad_outputs per_callable_static_grad_inputs[per_callable_bwd_idx] = static_grad_inputs + returned_param_grad_clone_slots = _returned_param_grad_clone_slots( + static_grad_inputs, + per_callable_module_params[per_callable_bwd_idx], + ) + per_callable_returned_param_grad_clone_slots[per_callable_bwd_idx] = ( + returned_param_grad_clone_slots + ) # Weak ref the static outputs and static grad inputs that are no longer needed # in the following steps. These two type of tensors are both in cudagraph @@ -727,6 +887,18 @@ def hook_fn( static_outputs ) + # Parameter grads can be weak-refed here only if they will be cloned + # before returning from Graphed.backward. + static_grad_inputs = per_callable_static_grad_inputs[per_callable_bwd_idx] + per_callable_static_grad_inputs[per_callable_bwd_idx] = tuple( + ( + make_weak_ref(grad_input) + if returned_param_grad_clone_slots[idx] and grad_input is not None + else grad_input + ) + for idx, grad_input in enumerate(static_grad_inputs) + ) + # Weak ref the static grad inputs of the previous backward pass within the # same chunk. if previous_per_callable_bwd_idx is not None: @@ -754,12 +926,16 @@ def hook_fn( # Capture forward graphs per_callable_static_outputs = [] per_callable_output_unflatten_spec = [] - graph_id = 0 - for func, args, kwargs, fwd_graph in zip(callables, sample_args, sample_kwargs, fwd_graphs): + for func_idx, (func, args, kwargs, fwd_graph) in enumerate( + zip(callables, sample_args, sample_kwargs, fwd_graphs) + ): + _run_capture_time_hooks(capture_time_hooks, func_idx, "forward_pre_hooks", func) + with _graph_context_wrapper(fwd_graph, stream=stream, pool=mempool): outputs = func(*args, **kwargs) - graph_callables[graph_id] = func - graph_id += 1 + graph_callables[func_idx] = func + + _run_capture_time_hooks(capture_time_hooks, func_idx, "forward_hooks", func) flatten_outputs, spec = _tree_flatten(outputs) per_callable_static_outputs.append(tuple(flatten_outputs)) @@ -768,6 +944,7 @@ def hook_fn( # Capture backward graphs in reverse order per_callable_static_grad_outputs = [] per_callable_static_grad_inputs = [] + per_callable_returned_param_grad_clone_slots = [] for static_input_surface, static_outputs, bwd_graph, bwd_dw_graph, bwd_idx in zip( reversed(per_callable_static_input_surfaces), reversed(per_callable_static_outputs), @@ -775,12 +952,19 @@ def hook_fn( reversed(bwd_dw_graphs), reversed(range(len(per_callable_static_input_surfaces))), ): - # For now, assumes all static_outputs require grad static_grad_outputs = tuple( torch.empty_like(o) if o is not None and o.requires_grad else None for o in static_outputs ) + grad_inputs: Tuple[Optional[torch.Tensor], ...] = () if is_training: + _run_capture_time_hooks( + capture_time_hooks, + bwd_idx, + "backward_pre_hooks", + callables[bwd_idx], + ) + inputs = tuple(i for i in static_input_surface if i.requires_grad) with _none_grad_context_wrapper(inputs), _graph_context_wrapper( bwd_graph, stream=stream, pool=mempool @@ -792,6 +976,13 @@ def hook_fn( ) grad_inputs = tuple(input.grad for input in inputs) + _run_capture_time_hooks( + capture_time_hooks, + bwd_idx, + "backward_hooks", + callables[bwd_idx], + ) + if need_bwd_dw_graph[bwd_idx]: with _graph_context_wrapper(bwd_dw_graph, stream=stream, pool=mempool): for module in visited_te_modules[bwd_idx]: @@ -812,10 +1003,19 @@ def hook_fn( per_callable_static_grad_outputs.append(static_grad_outputs) per_callable_static_grad_inputs.append(static_grad_inputs) + per_callable_returned_param_grad_clone_slots.append( + _returned_param_grad_clone_slots( + static_grad_inputs, + per_callable_module_params[bwd_idx], + ) + ) - # Reverses the most recent two lists + # Reverse the most recent per-callable lists. per_callable_static_grad_outputs = list(reversed(per_callable_static_grad_outputs)) per_callable_static_grad_inputs = list(reversed(per_callable_static_grad_inputs)) + per_callable_returned_param_grad_clone_slots = list( + reversed(per_callable_returned_param_grad_clone_slots) + ) # Now for every per_callable list, per_callable_*[i] holds the stuff for the ith callable. def make_graphed_autograd_function( @@ -829,6 +1029,7 @@ def make_graphed_autograd_function( static_outputs, static_grad_outputs, static_grad_inputs, + returned_param_grad_clone_slots, ): class Graphed(torch.autograd.Function): """Autograd function for graph replay.""" @@ -910,9 +1111,17 @@ def backward(ctx, *grads): "Expected static_grad_inputs to be a tuple, but got" f" {type(static_grad_inputs).__name__}" ) - return (None, None, None) + tuple( - b.detach() if b is not None else b for b in static_grad_inputs - ) + grad_inputs = [] + for idx, grad_input in enumerate(static_grad_inputs): + if grad_input is None: + grad_inputs.append(None) + elif returned_param_grad_clone_slots[idx]: + # Returned parameter grads may be installed directly as param.grad. + # Clone to avoid exposing CUDA graph static buffers to autograd users. + grad_inputs.append(grad_input.detach().clone()) + else: + grad_inputs.append(grad_input.detach()) + return (None, None, None) + tuple(grad_inputs) def functionalized(*user_args, **user_kwargs): @@ -1007,6 +1216,7 @@ def reset(): per_callable_static_outputs[i], per_callable_static_grad_outputs[i], per_callable_static_grad_inputs[i], + per_callable_returned_param_grad_clone_slots[i], ) func = graph_callables[i] @@ -1149,8 +1359,10 @@ def make_graphed_callables( pool: Optional[Tuple[int, ...]] = None, retain_graph_in_backward: bool = False, _reuse_graph_input_output_buffers: bool = False, + clone_param_grads_on_return: bool = True, pre_warmup_hook: Optional[Callable] = None, post_warmup_hook: Optional[Callable] = None, + capture_time_hooks: Optional[List[Optional[Dict[str, Dict]]]] = None, ) -> Union[Callable, Tuple[Callable, ...]]: """ Make CUDA graph version of Transformer Engine modules @@ -1189,10 +1401,34 @@ def make_graphed_callables( graphs. Only supported with Mcore interleaved pipeline parallelism, i.e. when `_order` is provided. All callables in `modules` are assumed to have inputs and outputs with the same dtype and shape. + clone_param_grads_on_return: bool, default = True + Clone parameter gradients before returning them from CUDA graph replay. + Disabling this avoids the extra clone/copy and may improve performance, + but returned parameter gradients will alias CUDA graph static gradient + buffers. These tensors no longer have standard PyTorch returned-gradient + lifetime semantics: a later replay of the same graph, or reused-buffer + replay of another callable, may overwrite retained hook or `.grad` + tensors. Only disable this when the caller consumes returned parameter + gradients before any such overwrite can occur. pre_warmup_hook: callable, default = None - A hook function that will be called before the warmup iterations. + A hook function that will be called once before all warmup iterations + (not once per callable). post_warmup_hook: callable, default = None - A hook function that will be called after the warmup iterations. + A hook function that will be called once after all warmup iterations + (not once per callable). + capture_time_hooks: list of dict, optional + Per-callable hooks invoked at capture time (during warmup iterations and + graph capture), but intentionally executed **outside** the CUDA graph + capture context so they are **not** recorded into the graph and will + **not** be replayed. Use this for operations that are inherently + non-capturable but essential for correct module execution, such as + CPU-side state updates. All hooks must have signature ``hook(module)`` + and must return ``None``. Any non-``None`` return value raises + ``RuntimeError``. + Each element corresponds to one callable and is a dict with any subset + of these keys: ``"forward_pre_hooks"``, ``"forward_hooks"``, + ``"backward_pre_hooks"``, and ``"backward_hooks"``. Each value is a + ``{hook_id: hook_fn}`` dict. Quantization parameters ----------------------- @@ -1323,6 +1559,12 @@ def make_graphed_callables( recipe = None module_uses_fp8 = dict(zip((id(m) for m in modules), enabled)) + # Canonicalize capture_time_hooks kwarg. + capture_time_hooks = _canonicalize_capture_time_hooks( + len(modules), + capture_time_hooks, + ) + # Store FP8 tensors to reset later. saved_fp8_tensors = save_fp8_tensors(modules, recipe=recipe) @@ -1384,8 +1626,10 @@ def call_func(self, *args, **kwargs): pool=pool, retain_graph_in_backward=retain_graph_in_backward, _reuse_graph_input_output_buffers=_reuse_graph_input_output_buffers, + clone_param_grads_on_return=clone_param_grads_on_return, pre_warmup_hook=pre_warmup_hook, post_warmup_hook=post_warmup_hook, + capture_time_hooks=capture_time_hooks, ) # Ensures warmup does not affect numerics for ops such as dropout. diff --git a/transformer_engine/pytorch/module/__init__.py b/transformer_engine/pytorch/module/__init__.py index 3cf15efc11..98ab745448 100644 --- a/transformer_engine/pytorch/module/__init__.py +++ b/transformer_engine/pytorch/module/__init__.py @@ -5,7 +5,7 @@ """Module level PyTorch APIs""" from .layernorm_linear import LayerNormLinear from .linear import Linear -from .grouped_linear import GroupedLinear +from .grouped_linear import GroupedLinear, is_module_grouped_tensor_path_supported from .layernorm_mlp import LayerNormMLP from .layernorm import LayerNorm from .rmsnorm import RMSNorm diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py index 1a4c933b6e..ad991474b8 100644 --- a/transformer_engine/pytorch/module/_common.py +++ b/transformer_engine/pytorch/module/_common.py @@ -17,6 +17,7 @@ from .. import cpp_extensions as tex from ..constants import TE_DType from ..export import is_in_onnx_export_mode +from ..tensor.hybrid_tensor import HybridQuantizer from ..utils import get_default_init_method if IS_HIP_EXTENSION: @@ -43,6 +44,41 @@ def set_quantizer_amax_reduction_group(quantizer, amax_reduction_group) -> None: target.amax_reduction_group = amax_reduction_group +def set_quantizer_usage_for_wgrad_all_gather(quantizer) -> None: + """Configure an all-gather output for consumption by wgrad.""" + if quantizer is None: + return + + parent_quantizer = getattr(quantizer, "parent_quantizer", None) + target = parent_quantizer if parent_quantizer is not None else quantizer + + # Hybrid currently gathers in high precision, then quantizes the full + # result, so request the columnwise representation consumed by wgrad. + if isinstance(target, HybridQuantizer): + rowwise_usage, columnwise_usage = False, True + elif quantizer.supports_only_rowwise_all_gather(): + # Per-tensor FP8 gathers rowwise data and synthesizes its transpose. + rowwise_usage, columnwise_usage = True, False + else: + rowwise_usage, columnwise_usage = False, True + + # Preserve wrapper-specific bookkeeping. In particular, DebugQuantizer + # propagates usage to its parent while keeping its own state synchronized. + quantizer.set_usage(rowwise=rowwise_usage, columnwise=columnwise_usage) + + +def can_reconstruct_wgrad_input_from_original(quantizer) -> bool: + """Whether wgrad input can be reconstructed from a saved original tensor.""" + target = getattr(quantizer, "parent_quantizer", quantizer) + if target is None: + target = quantizer + if isinstance(target, HybridQuantizer): + if target.columnwise_source == "original": + return True + return target.rowwise_quantizer.is_requantization_safe() + return target.is_requantization_safe() + + def _get_normalization_func(normalization: str, forward: bool): use_rmsnorm_triton = bool( int(os.environ.get('NVTE_USE_RMSNORM_TRITON', '0')) ) and IS_HIP_EXTENSION use_layernorm_triton = bool( int(os.environ.get('NVTE_USE_LAYERNORM_TRITON', '0')) ) and IS_HIP_EXTENSION diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index a8d4283b4e..3e124e1ea4 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -57,16 +57,20 @@ from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.nvfp4_tensor import NVFP4Quantizer from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer +from ..tensor.hybrid_tensor import HybridQuantizer +from ..tensor.identity_tensor import IdentityQuantizer if IS_HIP_EXTENSION: from ..tensor.fsdp2_allgather_tensor import FSDPAGTensor from ..triton_kernels.cast import te_quantize_triton from ..tensor.storage.float8_tensor_storage import Float8TensorStorage from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from ..tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage +from ..tensor.storage.hybrid_tensor_storage import HybridQuantizedTensorStorage from ..utils import ( get_device_compute_capability, is_non_tn_fp8_gemm_supported, torch_get_autocast_gpu_dtype, + get_device_compute_capability, get_nvtx_range_context, nvtx_range_push, nvtx_range_pop, @@ -98,6 +102,27 @@ layers_atomic_ring_exchange = [] +def _get_high_precision_init_val(parameter: torch.Tensor) -> Optional[torch.Tensor]: + """Return temporary pre-quantization initialization stored on a parameter.""" + return getattr(parameter, "_high_precision_init_val", None) + + +def _clear_high_precision_init_val(parameter: torch.Tensor) -> None: + """Release temporary pre-quantization initialization stored on a parameter.""" + if hasattr(parameter, "_high_precision_init_val"): + del parameter._high_precision_init_val + + +def _attach_high_precision_init_val( + parameter: torch.Tensor, + high_precision_init_val: torch.Tensor, +) -> None: + """Attach TE's temporary high-precision initialization contract to a parameter.""" + parameter._high_precision_init_val = high_precision_init_val + parameter.get_high_precision_init_val = MethodType(_get_high_precision_init_val, parameter) + parameter.clear_high_precision_init_val = MethodType(_clear_high_precision_init_val, parameter) + + def is_ub_initialized() -> bool: """Whether the Userbuffers communicators have been initialized.""" return _ub_initialized @@ -857,6 +882,14 @@ def _is_weight_workspace_valid( return False if quantizer.columnwise_usage and workspace._columnwise_data is None: return False + elif isinstance(workspace, HybridQuantizedTensorStorage): + # Workspace cached under one flag setting (e.g. inference with + # ``columnwise=False``) becomes stale when the next call needs the + # missing direction; invalidate so a fresh workspace is built. + if quantizer.rowwise_usage and workspace._rowwise_storage is None: + return False + if quantizer.columnwise_usage and workspace._columnwise_storage is None: + return False if isinstance(workspace, DebugQuantizedTensor) != isinstance(quantizer, DebugQuantizer): return False return True @@ -1013,6 +1046,35 @@ def module_setattr(self, name: str, value: Any) -> None: """ super().__setattr__(name, value) + def _apply(self, *args, **kwargs): + """Re-attach attributes that ``swap_tensors`` moves off a quantized parameter. + + ``_apply`` moves wrapper subclasses by exchanging the parameter's whole + ``__dict__``, which carries the inner buffers over but takes externally + attached state (``_high_precision_init_val``, ``main_grad``, ...) with it. + """ + snapshots = { + name: (param, dict(param.__dict__)) + for name, param in self._parameters.items() + if isinstance(param, QuantizedTensorStorage) + } + out = super()._apply(*args, **kwargs) + for name, (old_param, attrs) in snapshots.items(): + new_param = self._parameters.get(name) + if new_param is None: + raise RuntimeError( + f"{type(self).__name__}.{name} disappeared during _apply; the state" + " attached to it cannot be restored" + ) + for key, value in attrs.items(): + # Still present -> tensor state; the post-swap value is the right one. + if key in new_param.__dict__: + continue + if isinstance(value, MethodType) and value.__self__ is old_param: + value = MethodType(value.__func__, new_param) + setattr(new_param, key, value) + return out + @property def output_quantizer_role(self) -> Optional[QuantizerRole]: """Caller-configurable :class:`QuantizerRole` for the forward output quantizer. @@ -1156,6 +1218,8 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: # Return early if recipe state matches recipe if self.fp8_meta_tensors_initialized: recipe_state = self.fp8_meta[fp8_meta_tensor_key] + # TODO(#3157): Match built-in recipes by full config, not just RecipeState type, so + # same-class mid-training changes rebuild quantizers/workspaces correctly. if recipe.delayed() and isinstance(recipe_state, DelayedScalingRecipeState): self.adjust_amax_history_length(recipe.amax_history_len, fwd=fwd) return @@ -1174,6 +1238,9 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: if recipe.nvfp4() and isinstance(recipe_state, NVFP4BlockScalingRecipeState): return if recipe.custom() and isinstance(recipe_state, CustomRecipeState): + # TODO(#3157): Compare CustomRecipe/qfactory config here. qfactory changes made + # mid-training on the same recipe object currently do not take effect because + # stale quantizers are reused. if recipe_state.recipe is recipe: return @@ -1317,6 +1384,40 @@ def _get_weight_quantizers(self) -> List[Quantizer]: f"{self.__class__.__name__} class does not implement _get_weight_quantizers function" ) + def _enable_weight_preswizzle( + self, + quantizer: Quantizer, + weight: torch.Tensor, + ) -> bool: + """Whether to fuse scale-factor swizzling into weight quantization. + + When enabled, scales are preswizzled during quantization instead of lazily + inside every GEMM. Disabled when primary weights are already quantized + (dequant and FSDP2 all-gather expect the unswizzled layout). For NVFP4, + enabled only for shapes/architectures where the fused swizzle+quantize + kernel is supported. Weight quantization always uses the single-tensor + kernel (including GroupedLinear's cached weights), so NVFP4 RHT + eligibility uses that kernel's 64-row alignment. + """ + if self.primary_weights_in_fp8: + return False + if isinstance(quantizer, MXFP8Quantizer): + return True + if isinstance(quantizer, NVFP4Quantizer): + rows, cols = weight.numel() // weight.shape[-1], weight.shape[-1] + arch_supported = get_device_compute_capability() >= (10, 0) + if quantizer.with_rht: + return arch_supported and rows % 64 == 0 and cols % 128 == 0 + return ( + arch_supported + and quantizer.with_2d_quantization + and not quantizer.row_scaled_nvfp4 + and not quantizer.nvfp4_use_4over6 + and rows % 128 == 0 + and cols % 128 == 0 + ) + return False + def init_fp8_meta_tensors(self, recipe: Recipe) -> None: """Init scales and amaxes.""" self.set_meta_tensor(True, recipe) @@ -1794,8 +1895,12 @@ def grad_output_preprocess( ): grad_bias = grad_output.dequantize().view(-1, grad_output.shape[-1]).sum(dim=0) else: - if isinstance(quantizer, Float8BlockQuantizer): - # unfuse bgrad for now until cast_transpose + dgrad calculation is ready for Float8BlockQuantizer. + if isinstance( + quantizer, (Float8BlockQuantizer, HybridQuantizer, IdentityQuantizer) + ): + # Float8BlockQuantizer: unfused until cast_transpose + dgrad is ready. + # HybridQuantizer: tex.bgrad_quantize doesn't recognize hybrid quantizers. + # IdentityQuantizer: high-precision passthrough; bgrad computed in HP. grad_bias = grad_output.view(-1, grad_output.shape[-1]).sum(dim=0) else: grad_bias, grad_output = tex.bgrad_quantize(grad_output, quantizer) @@ -1865,8 +1970,11 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: if IS_HIP_EXTENSION and not self.keep_fp8_weight_transpose_cache: quantizer.columnwise_usage=False + # HybridQuantizer is included so its current-scaling / NVFP4 + # sub-quantizers get the same cross-shard amax reduction as the + # vanilla path (no-op for block-scaled sub-quantizers like MXFP8). if is_dtensor and isinstance( - quantizer, (Float8CurrentScalingQuantizer, NVFP4Quantizer) + quantizer, (Float8CurrentScalingQuantizer, NVFP4Quantizer, HybridQuantizer) ): device_mesh = dtensor_param.device_mesh amax_reduction_group = ( @@ -1918,21 +2026,10 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: # should call `clear_high_precision_init_val` to remove it after master weight # is initialized. - def get(self): - if hasattr(self, "_high_precision_init_val"): - return self._high_precision_init_val - return None - - def clear(self): - if hasattr(self, "_high_precision_init_val"): - del self._high_precision_init_val - # DTensor.from_local() does not preserve object identity, # so attach to the DTensor's local tensor when applicable. target = dtensor_param._local_tensor if is_dtensor else param - target._high_precision_init_val = high_precision_init_val - target.get_high_precision_init_val = MethodType(get, target) - target.clear_high_precision_init_val = MethodType(clear, target) + _attach_high_precision_init_val(target, high_precision_init_val) if not is_dtensor: self.module_setattr(name, param) diff --git a/transformer_engine/pytorch/module/fp8_padding.py b/transformer_engine/pytorch/module/fp8_padding.py index 3a0073a492..bb2cfc7dd8 100644 --- a/transformer_engine/pytorch/module/fp8_padding.py +++ b/transformer_engine/pytorch/module/fp8_padding.py @@ -82,9 +82,10 @@ class Fp8Padding(torch.nn.Module): num_gemms : int number of GEMMs to be performed simultaneously. align_size : int, optional - the alignment size for the input tensor. If not provided, the alignment size will - be determined by the FP8/FP4 recipe (32 for MXFP8/NVFP4 and 16 for others) in the first - forward pass. + Alignment size for each grouped input. If not provided, it is + determined from the active recipe on the first forward pass: + 32 for MXFP8, 128 for NVFP4, 16 for other built-in recipes, + and ``CustomRecipe.quantization_alignment`` for custom recipes. TODO: invesitgate the alignment requirement for non-mxfp8 cases on ROCm """ diff --git a/transformer_engine/pytorch/module/fp8_unpadding.py b/transformer_engine/pytorch/module/fp8_unpadding.py index c5d396837f..d34605e195 100644 --- a/transformer_engine/pytorch/module/fp8_unpadding.py +++ b/transformer_engine/pytorch/module/fp8_unpadding.py @@ -78,9 +78,10 @@ class Fp8Unpadding(torch.nn.Module): num_gemms : int number of GEMMs to be performed simultaneously. align_size : int, optional - The alignment size for the input tensor. If not provided, the alignment size will - be automatically determined based on the FP8/FP4 recipe in the first forward pass: - 32 for MXFP8 or NVFP4, otherwise 16. + Alignment size used to pad each grouped input. If not provided, + it is determined from the active recipe on the first forward + pass: 32 for MXFP8, 128 for NVFP4, 16 for other built-in + recipes, and ``CustomRecipe.quantization_alignment`` for custom recipes. """ def __init__( diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 910a726b40..7098a93792 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -6,7 +6,7 @@ """GroupedLinear API""" -from typing import Union, Optional, Callable, Tuple, List +from typing import Union, Optional, Callable, Tuple, List, Sequence from itertools import chain import os import warnings @@ -29,8 +29,11 @@ _2X_ACC_FPROP, _2X_ACC_DGRAD, _2X_ACC_WGRAD, + _attach_high_precision_init_val, + _clear_high_precision_init_val, + _get_high_precision_init_val, ) -from ._common import WeightGradStore +from ._common import can_reconstruct_wgrad_input_from_original, WeightGradStore from ..quantization import FP8GlobalStateManager, QuantizerRole from ..utils import ( divide, @@ -48,6 +51,12 @@ is_fp8_activation_recompute_enabled, in_fp8_activation_recompute_phase, ) +from ..distributed_weight import ( + is_distributed_weight, + materialize_weight_for_forward, + materialize_weight_for_backward, + finalize_weight_grads, +) from ..cpp_extensions import ( general_grouped_gemm, general_grouped_gemm_for_grouped_tensor, @@ -57,7 +66,15 @@ from ..cpu_offload import is_cpu_offload_enabled, mark_not_offload, start_offload from ..triton.grouped_dbias_dscales import compute_grouped_dbias -from ..tensor import Float8CurrentScalingQuantizer, Float8Quantizer, MXFP8Quantizer, NVFP4Quantizer +from ..tensor import ( + Float8BlockQuantizer, + Float8CurrentScalingQuantizer, + Float8Quantizer, + HybridQuantizer, + IdentityQuantizer, + MXFP8Quantizer, + NVFP4Quantizer, +) from ..quantized_tensor import ( QuantizedTensorStorage, Quantizer, @@ -72,7 +89,436 @@ from transformer_engine.pytorch.triton_kernels.grouped_gemm import general_grouped_gemm_triton import os -__all__ = ["GroupedLinear"] +_NATIVE_SPLIT_QUANTIZER_TYPES = frozenset( + { + Float8Quantizer, + Float8CurrentScalingQuantizer, + Float8BlockQuantizer, + MXFP8Quantizer, + NVFP4Quantizer, + } +) + + +def _supports_native_split_quantize(quantizer): + """Whether ``tex.split_quantize`` has an exact converter for this quantizer.""" + return type(quantizer) in _NATIVE_SPLIT_QUANTIZER_TYPES + + +def _uses_identity_quantizer(quantizer): + """Whether a quantizer, including a hybrid sub-quantizer, is Identity-backed.""" + if quantizer is None: + return False + if isinstance(quantizer, IdentityQuantizer): + return True + if isinstance(quantizer, HybridQuantizer): + return _uses_identity_quantizer(quantizer.rowwise_quantizer) or _uses_identity_quantizer( + quantizer.columnwise_quantizer + ) + return False + + +def _identity_quantizer_signature(quantizer): + """Identity usage per GEMM direction: (rowwise, columnwise).""" + if isinstance(quantizer, HybridQuantizer): + return ( + _uses_identity_quantizer(quantizer.rowwise_quantizer), + _uses_identity_quantizer(quantizer.columnwise_quantizer), + ) + identity = isinstance(quantizer, IdentityQuantizer) + return (identity, identity) + + +_DYNAMIC_QUANTIZER_SIGNATURE_FIELDS = frozenset( + { + "rowwise_usage", + "columnwise_usage", + "internal", + "optimize_for_gemm", + } +) + + +def _backend_quantizer_signature(quantizer): + """Return backend configuration that grouped kernels require to be uniform.""" + if quantizer is None: + return None + + # Identity is not registered as a torch.compile value quantizer, but its + # dtype changes the grouped GEMM input type and therefore must be uniform. + if isinstance(quantizer, IdentityQuantizer): + return (type(quantizer), (("dtype", quantizer.dtype),)) + + fields = quantizer._value_fields() + if fields is None: + # Delayed-scaling Float8Quantizer carries per-expert scale/amax tensors, + # which are intentionally different, but its emitted FP8 dtype is a + # group-wide backend choice. Other unregistered/custom quantizers retain + # the conservative exact-family behavior until they expose value fields. + fields = ("dtype",) if isinstance(quantizer, Float8Quantizer) else () + + config = [] + for name in fields: + if name in _DYNAMIC_QUANTIZER_SIGNATURE_FIELDS: + continue + value = getattr(quantizer, name) + if name == "dtype": + value = int(value) + config.append((name, value)) + return (type(quantizer), tuple(config)) + + +def _validate_backend_match(reference, quantizer, operand_name, direction, expert_index): + """Validate one expert against the group's reference backend.""" + if type(quantizer) is not type(reference): + raise ValueError( + f"GroupedLinear {operand_name} quantizers use incompatible {direction} backend" + f" families across experts: expert 0 uses {type(reference).__name__}, but expert" + f" {expert_index} uses {type(quantizer).__name__}. Grouped operands require one" + " quantizer family per direction." + ) + reference_signature = _backend_quantizer_signature(reference) + quantizer_signature = _backend_quantizer_signature(quantizer) + if quantizer_signature != reference_signature: + raise ValueError( + f"GroupedLinear {operand_name} quantizers use incompatible {direction} backend" + f" configurations across experts: expert 0 uses {reference_signature}, but expert" + f" {expert_index} uses {quantizer_signature}. Grouped operands require the same" + " backend-relevant configuration per direction." + ) + + +def _validate_grouped_quantizer_list(quantizers, *, operand_name="operand") -> None: + """Validate one grouped operand once when its quantizer generation changes.""" + if not quantizers: + return + + reference = quantizers[0] + reference_is_hybrid = isinstance(reference, HybridQuantizer) + reference_identity = _identity_quantizer_signature(reference) + + for expert_index, quantizer in enumerate(quantizers[1:], start=1): + if (quantizer is None) != (reference is None): + raise ValueError( + f"GroupedLinear {operand_name} quantizers mix None and concrete quantizers" + f" across experts: expert 0 is {type(reference).__name__}, but expert" + f" {expert_index} is {type(quantizer).__name__}." + ) + if reference is None: + continue + + quantizer_is_hybrid = isinstance(quantizer, HybridQuantizer) + if quantizer_is_hybrid != reference_is_hybrid: + raise ValueError( + f"GroupedLinear {operand_name} quantizers mix HybridQuantizer and non-hybrid" + f" quantizers across experts: expert 0 is {type(reference).__name__}, but expert" + f" {expert_index} is {type(quantizer).__name__}." + ) + + identity = _identity_quantizer_signature(quantizer) + if identity != reference_identity: + raise ValueError( + f"GroupedLinear {operand_name} quantizers mix Identity-backed and quantized" + f" directions across experts: expert 0 uses {reference_identity}, but expert" + f" {expert_index} uses {identity}." + ) + + if reference_is_hybrid: + _validate_backend_match( + reference.rowwise_quantizer, + quantizer.rowwise_quantizer, + operand_name, + "rowwise", + expert_index, + ) + _validate_backend_match( + reference.columnwise_quantizer, + quantizer.columnwise_quantizer, + operand_name, + "columnwise", + expert_index, + ) + if quantizer.columnwise_source != reference.columnwise_source: + raise ValueError( + f"GroupedLinear {operand_name} HybridQuantizer list has mixed columnwise" + " source policies across experts: expert 0 uses" + f" {reference.columnwise_source!r}, but expert {expert_index} uses" + f" {quantizer.columnwise_source!r}." + ) + else: + _validate_backend_match( + reference, + quantizer, + operand_name, + "plain", + expert_index, + ) + + +def _split_quantize_non_hybrid( + tensor, + m_splits, + quantizers, + activation_dtype, + *, + disable_bulk_allocation=False, + allow_identity_views=True, +): + """Split and quantize one homogeneous, non-Hybrid quantizer list.""" + reference = quantizers[0] + if _supports_native_split_quantize(reference): + return tex.split_quantize( + tensor, + m_splits, + quantizers, + disable_bulk_allocation=disable_bulk_allocation, + ) + + tensor = cast_if_needed(tensor, activation_dtype) + if ( + allow_identity_views + # Only the base IdentityQuantizer can bypass quantization; subclasses + # may override its behavior and must go through their normal call path. + and type(reference) is IdentityQuantizer # pylint: disable=unidiomatic-typecheck + and (reference.dtype is None or reference.dtype == activation_dtype) + ): + return torch.split(tensor, m_splits) + + return [ + quantizer(tensor_part) if quantizer is not None else tensor_part + for tensor_part, quantizer in zip(torch.split(tensor, m_splits), quantizers) + ] + + +def _split_quantize_hybrid( + tensor, + m_splits, + quantizers, + *, + disable_bulk_allocation=False, +): + """Grouped split+quantize for an all-hybrid, generation-validated operand.""" + from ..tensor.storage.hybrid_tensor_storage import HybridQuantizedTensorStorage as HybridStorage + + reference = quantizers[0] + rowwise_enabled = reference.rowwise_usage + columnwise_enabled = reference.columnwise_usage + columnwise_source = reference.columnwise_source + rowwise_quantizers = [quantizer.rowwise_quantizer for quantizer in quantizers] + columnwise_quantizers = [quantizer.columnwise_quantizer for quantizer in quantizers] + + needs_rowwise_result = rowwise_enabled or ( + columnwise_enabled and columnwise_source == "rowwise_dequantized" + ) + row_results = ( + _split_quantize_non_hybrid( + tensor, + m_splits, + rowwise_quantizers, + tensor.dtype, + disable_bulk_allocation=disable_bulk_allocation, + allow_identity_views=False, + ) + if needs_rowwise_result + else [None] * len(quantizers) + ) + + columnwise_src = tensor + if columnwise_enabled and columnwise_source == "rowwise_dequantized": + # Assemble the exact grouped row results in split order. NVFP4 padding + # and scale layout can differ from independently quantizing each split. + columnwise_src = torch.cat( + [result.dequantize(dtype=tensor.dtype) for result in row_results], + dim=0, + ) + col_results = ( + _split_quantize_non_hybrid( + columnwise_src, + m_splits, + columnwise_quantizers, + tensor.dtype, + disable_bulk_allocation=disable_bulk_allocation, + allow_identity_views=False, + ) + if columnwise_enabled + else [None] * len(quantizers) + ) + + return [ + HybridStorage( + rowwise_storage=row if rowwise_enabled else None, + columnwise_storage=col, + quantizer=q, + fake_dtype=tensor.dtype, + ) + for row, col, q in zip( + row_results, + col_results, + quantizers, + ) + ] + + +def _split_quantize( + tensor: torch.Tensor, + split_sizes: List[int], + with_quantized_output: bool, + quantizers: Optional[List[Quantizer]], + dtype: torch.dtype, + with_debug_quantizers: bool, + disable_bulk_allocation: bool, +) -> Sequence[Union[torch.Tensor, QuantizedTensorStorage]]: + """Split a tensor and quantize each part if needed.""" + if not with_quantized_output: + return torch.split(cast_if_needed(tensor, dtype), split_sizes) + + if quantizers is None or quantizers[0] is None: + raise ValueError("Quantizers are required for quantized split output") + + if with_debug_quantizers: + return DebugQuantizer.multi_tensor_quantize(tensor, quantizers, split_sizes, dtype) + + reference = quantizers[0] + if isinstance(reference, HybridQuantizer): + return _split_quantize_hybrid( + tensor, + split_sizes, + quantizers, + disable_bulk_allocation=disable_bulk_allocation, + ) + + return _split_quantize_non_hybrid( + tensor, + split_sizes, + quantizers, + dtype, + disable_bulk_allocation=disable_bulk_allocation, + ) + + +def _split_quantize_and_bias( + tensor: torch.Tensor, + split_sizes: List[int], + *, + fp8: bool, + debug: bool, + quantizers: Optional[List[Quantizer]], + dtype: torch.dtype, + use_bias: bool, + recipe: Recipe, + disable_bulk_allocation: bool, +) -> Tuple[ + Sequence[Union[torch.Tensor, QuantizedTensorStorage]], + List[Optional[torch.Tensor]], +]: + """Split grad output, quantize if needed, and compute unfused bias gradients.""" + num_splits = len(split_sizes) + grad_biases = [None] * num_splits + reference = quantizers[0] + identity = _uses_identity_quantizer(reference) + hybrid = isinstance(reference, HybridQuantizer) and not identity + + use_native_bgrad_quantize = ( + fp8 + and not debug + and not hybrid + and use_bias + and not identity + and (recipe.delayed() or recipe.float8_current_scaling() or recipe.mxfp8()) + ) + if use_native_bgrad_quantize: + outputs = [None] * num_splits + for i, tensor_part in enumerate(torch.split(tensor, split_sizes)): + grad_biases[i], outputs[i] = tex.bgrad_quantize(tensor_part, quantizers[i]) + return outputs, grad_biases + + with_quantized_output = fp8 or debug + if with_quantized_output and (use_bias or debug): + for i, tensor_part in enumerate(torch.split(tensor, split_sizes)): + grad_biases[i] = tensor_part.sum(dim=0) + + # Preserve the existing CPU-offload policy: only Hybrid split-quantize + # disables bulk allocation in backward. + disable_bulk_allocation = disable_bulk_allocation if hybrid else False + outputs = _split_quantize( + tensor, + split_sizes, + with_quantized_output=with_quantized_output, + quantizers=quantizers, + dtype=dtype, + with_debug_quantizers=debug, + disable_bulk_allocation=disable_bulk_allocation, + ) + return outputs, grad_biases + + +__all__ = ["GroupedLinear", "is_module_grouped_tensor_path_supported"] + + +def is_module_grouped_tensor_path_supported( + recipe: Optional[Recipe], + dtype: torch.dtype, +) -> bool: + """Whether the module grouped-tensor path supports this recipe and dtype. + + The grouped-tensor path dispatches to ``general_grouped_gemm_for_grouped_tensor`` + and does not inspect split values because they may reside in a CUDA tensor. + Inspecting them on the host would add synchronization and break CUDA Graph safety. + + Supported Compute Capability (CC) and precisions: + + * Hopper (CC 9.0): BF16/FP16, FP8 per-tensor current scaling, and FP8 + block scaling. + * Blackwell (CC 10.x and 11.0): BF16/FP16, FP8 per-tensor current scaling, + MXFP8, and NVFP4 with RHT. + * Custom recipes are unsupported because they may assign different + quantizers to input, weight, and grad-output roles. This predicate + currently supports only built-in recipes with known uniform layouts. + * FP8 delayed scaling is unsupported because the required grouped + quantization kernels are unavailable. + * FP8 block scaling is unsupported by this path on Blackwell because it + does not implement the legacy path's MXFP8-broadcast emulation. + * Grouped GEMM requires cuBLASLt 13.3+, with 13.4+ required on Hopper, + 13.5+ required for FP8 per-tensor current scaling on Hopper, and 13.6+ + required for FP8 block scaling on Hopper. + * FP32 is unsupported by the cuBLASLt grouped GEMM. + + Runtime-only restrictions such as debug mode, CPU offloading, calibration, + output quantization, and backend selection are checked separately by + ``GroupedLinear``. + """ + if dtype not in (torch.bfloat16, torch.float16): + return False + + device_capability = get_device_compute_capability() + if not (9, 0) <= device_capability <= (11, 0): + return False + cublaslt_version = tex.get_cublasLt_version() + if cublaslt_version < 130300: + return False + if device_capability < (10, 0) and cublaslt_version < 130400: + return False + + if recipe is None: + return True + if recipe.custom(): + return False + if recipe.backward_override is not None: + return False + if recipe.float8_current_scaling(): + return device_capability >= (10, 0) or cublaslt_version >= 130500 + if recipe.float8_block_scaling(): + # cuBLASLt 13.6 fixes Hopper grouped GEMM algo selection for block-scaled FP8. + return device_capability < (10, 0) and cublaslt_version >= 130600 + if recipe.mxfp8(): + return device_capability >= (10, 0) + if recipe.nvfp4(): + return ( + device_capability >= (10, 0) + and not recipe.disable_rht + and not recipe.row_scaled_activation + ) + return False class _GroupedLinear(torch.autograd.Function): @@ -90,88 +536,13 @@ def _maybe_dequantize( return tensor.dequantize(dtype=dtype) return cast_if_needed(tensor, dtype) - @staticmethod - def _is_grouped_tensor_path_supported( - *, - fp8: bool, - fp8_calibration: bool, - debug: bool, - cpu_offloading: bool, - backward_override: Optional[str], - save_original_input: bool, - activation_dtype: torch.dtype, - input_quantizers: List[Optional[Quantizer]], - output_quantizers: List[Optional[Quantizer]], - ) -> bool: - """Whether to use cuBLASLt grouped GEMM through GroupedTensor metadata. - - There are no checks whether split sizes are supported. Splits - may be in a CUDA tensor, so checking would hurt performance - and be incompatible with CUDA Graphs. - - Supported Compute Capability (CC) and precisions: - * Hopper (CC 9.0): BF16/FP16 and FP8 per-tensor current scaling. - * Blackwell (CC 10.x and 11.0): BF16/FP16/MXFP8/NVFP4 with RHT and FP8 - per-tensor current scaling. - FP8 delayed scaling and FP8 block scaling are not supported because the - corresponding grouped quantization kernels are missing. - Grouped GEMM requires cuBLAS 13.3+ (13.4+ on Hopper, 13.5+ for FP8 - per-tensor current scaling on Hopper); otherwise the legacy path is used. - Non-RHT NVFP4 falls back to the legacy path because graph-safe grouped quantization - currently requires RHT. - - Input/weight/grad_output quantizers are assumed to be of the same type, otherwise it would - trigger a fatal error in the cuBLASLt grouped GEMM check. - """ - # CUDA-only path; ROCm uses general_grouped_gemm. - if IS_HIP_EXTENSION: - return False - # 1. Filter by environment variable - if not bool(int(os.getenv("NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM", "0"))): - return False - # 2. Filter out advanced features - if ( - debug - or cpu_offloading - or fp8_calibration - or backward_override is not None - or save_original_input - ): - return False - # 3. Filter by compute capability and cuBLAS version - device_capability = get_device_compute_capability() - if not (9, 0) <= device_capability <= (11, 0): - return False - cublaslt_version = tex.get_cublasLt_version() - if cublaslt_version < 130300: - return False - if device_capability < (10, 0) and cublaslt_version < 130400: - return False - # 4. Output quantization is not supported. - if any(q is not None for q in output_quantizers): - return False - # 5. Filter by quantization recipes. - if fp8: - if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers): - # FP8 per-tensor scaling grouped GEMM on Hopper requires cuBLAS 13.5+. - if device_capability < (10, 0) and cublaslt_version < 130500: - return False - return True - # MXFP8 and NVFP4 require Blackwell+. - if not (10, 0) <= device_capability <= (11, 0): - return False - return all(isinstance(q, MXFP8Quantizer) for q in input_quantizers) or all( - isinstance(q, NVFP4Quantizer) and q.with_rht for q in input_quantizers - ) - return activation_dtype in (torch.bfloat16, torch.float16) - @staticmethod def _make_grouped_tensor( data: torch.Tensor, *, num_gemms: int, split_sizes: torch.Tensor, - base_split_offsets: torch.Tensor, + tensor_offsets: torch.Tensor, last_dim: int, dtype: torch.dtype, ) -> GroupedTensorStorage: @@ -183,7 +554,7 @@ def _make_grouped_tensor( quantizer=None, data=data.reshape(-1), first_dims=split_sizes, - tensor_offsets=base_split_offsets * last_dim, + tensor_offsets=tensor_offsets, ) @staticmethod @@ -214,14 +585,101 @@ def _prepare_weights_for_grouped_tensor_gemm( weight_quantizers: List[Optional[Quantizer]], weight_workspaces: List[Optional[QuantizedTensorStorage]], *, + num_gemms: int, + single_grouped_weight: bool, with_quantized_compute: bool, columnwise_usage: bool, activation_dtype: torch.dtype, is_first_microbatch: Optional[bool], skip_fp8_weight_update: Optional[torch.Tensor], cache_weight: bool, - ) -> Tuple[List[torch.Tensor], List[Optional[QuantizedTensorStorage]]]: - """Prepare discrete weight tensors for GroupedTensor GEMM.""" + ) -> Tuple[ + Union[GroupedTensorStorage, List[torch.Tensor]], + List[Optional[QuantizedTensorStorage]], + ]: + """Prepare a grouped parameter or discrete weights for GroupedTensor GEMM.""" + if single_grouped_weight: + weight = weights[0] + if not isinstance(weight, GroupedTensorStorage): + raise TypeError( + "single_grouped_weight requires the weight parameter to be a GroupedTensor." + ) + + new_workspaces: List[Optional[QuantizedTensorStorage]] = [None] + if weight.quantizer is not None: + if not with_quantized_compute: + raise RuntimeError( + "Quantized single grouped weights require quantized grouped GEMM compute." + ) + return weight, new_workspaces + + if not with_quantized_compute: + if weight.rowwise_data is None: + raise RuntimeError("Single grouped weight has no rowwise storage.") + if weight.rowwise_data.dtype == activation_dtype: + return weight, new_workspaces + data = weight.rowwise_data.to(dtype=activation_dtype) + return ( + GroupedTensorStorage( + shape=weight.logical_shape, + dtype=activation_dtype, + num_tensors=num_gemms, + shapes=weight.tensor_shapes, + quantizer=None, + data=data, + ), + new_workspaces, + ) + + weight_quantizer = weight_quantizers[0] + if weight_quantizer is None: + raise RuntimeError("Quantized grouped compute requires a weight quantizer.") + weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + # forward() already applied _enable_weight_preswizzle(); preserve that decision + # because not every quantizer and weight shape supports fused quantize-swizzle. + + workspace = weight_workspaces[0] if weight_workspaces else None + + if workspace is not None and ( + workspace.quantizer is not weight_quantizer + or (columnwise_usage and workspace.columnwise_data is None) + ): + workspace = None + + if weight.rowwise_data is None: + raise RuntimeError("Single grouped weight has no rowwise storage to quantize.") + source = weight.rowwise_data.view(weight.logical_shape) + update_workspace = is_first_microbatch is None or is_first_microbatch + if workspace is None: + if cache_weight: + # Match quantize_weight(): persistent workspaces must be Tensor subclasses + # so autograd can save them without decomposing their storage metadata. + saved_internal = weight_quantizer.internal + weight_quantizer.internal = False + grouped_weight = tex.group_quantize( + source, + weight_quantizer, + num_gemms, + None, + ) + if cache_weight: + weight_quantizer.internal = saved_internal + elif skip_fp8_weight_update is not None or update_workspace: + grouped_weight = tex.group_quantize( + source, + weight_quantizer, + num_gemms, + None, + noop_flag=skip_fp8_weight_update, + output=workspace, + ) + else: + grouped_weight = workspace + + if cache_weight: + new_workspaces[0] = grouped_weight + return grouped_weight, new_workspaces + weights_for_gemm: List[torch.Tensor] = [] new_workspaces: List[Optional[QuantizedTensorStorage]] = [None] * len(weights) if not with_quantized_compute: @@ -261,13 +719,10 @@ def _validate_or_alloc_output( """ if buffer is None: return torch.empty((rows, cols), dtype=dtype, device=device) - if buffer.dim() != 2: - raise ValueError(f"Output buffer must be 2D, got {buffer.dim()}D.") - if buffer.size(0) != rows: - raise ValueError(f"Output buffer rows {buffer.size(0)} must match input rows {rows}.") - if buffer.size(1) != cols: + expected_shape = (rows, cols) + if buffer.shape != expected_shape: raise ValueError( - f"Output buffer last dim {buffer.size(1)} does not match required {cols}." + f"Output buffer shape {tuple(buffer.shape)} must match required {expected_shape}." ) if buffer.dtype != dtype: raise ValueError(f"Output buffer dtype {buffer.dtype} does not match required {dtype}.") @@ -281,6 +736,43 @@ def _validate_or_alloc_output( raise ValueError("Output buffer must not require gradient.") return buffer + @staticmethod + def _prepare_bias_for_grouped_tensor_gemm( + biases: Tuple[torch.Tensor, ...], + *, + single_grouped_bias: bool, + num_gemms: int, + out_features: int, + dtype: torch.dtype, + ) -> GroupedTensorStorage: + """Prepare grouped or discrete bias storage for grouped GEMM.""" + if not single_grouped_bias: + return _GroupedLinear._make_grouped_bias( + biases, + num_gemms=num_gemms, + out_features=out_features, + dtype=dtype, + ) + + bias = biases[0] + if not isinstance(bias, GroupedTensorStorage): + raise TypeError("single_grouped_bias requires a GroupedTensor parameter.") + bias_data = bias.rowwise_data + if bias_data.dtype != dtype: + bias_data = bias_data.to(dtype=dtype) + + # The parameter exposes a packed [num_gemms, out_features] tensor, but its grouped + # members are 1D vectors. The grouped bias-add kernel consumes those same bytes as + # num_gemms row matrices with shape [1, out_features]. + return GroupedTensorStorage( + shape=(num_gemms, out_features), + dtype=dtype, + num_tensors=num_gemms, + shapes=[(1, out_features)] * num_gemms, + quantizer=None, + data=bias_data.reshape(-1), + ) + @staticmethod def _forward_grouped_tensor( ctx, @@ -302,6 +794,9 @@ def _forward_grouped_tensor( weight_workspaces: List[Optional[QuantizedTensorStorage]], cache_weight: bool, skip_fp8_weight_update: Optional[torch.Tensor], + save_original_input: bool, + single_grouped_weight: bool, + single_grouped_bias: bool, weights: Tuple[torch.Tensor, ...], biases: Tuple[torch.Tensor, ...], out: Optional[torch.Tensor] = None, @@ -311,11 +806,22 @@ def _forward_grouped_tensor( num_gemms = len(m_splits) device = inp.device in_features = weights[0].size(-1) - out_features = weights[0].size(0) + out_features = weights[0].size(-2) weight_requires_grad = weights[0].requires_grad - - split_sizes = m_splits.to(device=device) - base_split_offsets = tex.splits_to_offsets(split_sizes, 1) + save_original_input = save_original_input and weight_requires_grad + + split_sizes, ( + base_split_offsets, + input_tensor_offsets, + output_tensor_offsets, + ) = tex.splits_to_offsets_multi( + m_splits, + device, + strides=[1, in_features, out_features], + include_leading_zero=[True, True, True], + dtypes=[torch.int64, torch.int64, torch.int64], + bulk_allocate=True, + ) inp_view = inp.reshape(-1, in_features) x = cast_if_needed(inp_view, activation_dtype) @@ -323,16 +829,22 @@ def _forward_grouped_tensor( input_quantizer = input_quantizers[0] input_quantizer.set_usage( rowwise=True, - columnwise=is_grad_enabled and weight_requires_grad, + columnwise=(is_grad_enabled and weight_requires_grad and not save_original_input), ) input_quantizer.optimize_for_gemm = True - grouped_x = tex.group_quantize(x, input_quantizer, num_gemms, split_sizes) + grouped_x = tex.group_quantize( + x, + input_quantizer, + num_gemms, + split_sizes, + tensor_offsets=input_tensor_offsets, + ) else: grouped_x = _GroupedLinear._make_grouped_tensor( x, num_gemms=num_gemms, split_sizes=split_sizes, - base_split_offsets=base_split_offsets, + tensor_offsets=input_tensor_offsets, last_dim=in_features, dtype=activation_dtype, ) @@ -342,6 +854,8 @@ def _forward_grouped_tensor( weights, weight_quantizers, weight_workspaces, + num_gemms=num_gemms, + single_grouped_weight=single_grouped_weight, with_quantized_compute=fp8, columnwise_usage=columnwise_usage, activation_dtype=activation_dtype, @@ -361,15 +875,16 @@ def _forward_grouped_tensor( out, num_gemms=num_gemms, split_sizes=split_sizes, - base_split_offsets=base_split_offsets, + tensor_offsets=output_tensor_offsets, last_dim=out_features, dtype=activation_dtype, ) grouped_bias = None if use_bias: - grouped_bias = _GroupedLinear._make_grouped_bias( + grouped_bias = _GroupedLinear._prepare_bias_for_grouped_tensor_gemm( biases, + single_grouped_bias=single_grouped_bias, num_gemms=num_gemms, out_features=out_features, dtype=activation_dtype, @@ -391,26 +906,35 @@ def _forward_grouped_tensor( ) if is_grad_enabled: + input_to_save = grouped_x if weight_requires_grad: - # (For FP8 per tensor current scaling on Hopper --> Free Rowwise Data - # in backward pass) - if fp8 and grouped_x.columnwise_data is not None: + if save_original_input: + # Save the high-precision input and reconstruct the grouped columnwise + # operand in backward instead of retaining a second quantized copy. + input_to_save = inp + elif fp8 and grouped_x.columnwise_data is not None: + # Wgrad only consumes the columnwise representation. grouped_x.rowwise_data = None grouped_x.scale_inv = None else: - grouped_x = None + input_to_save = None + + weights_to_save = [weights_for_gemm] if single_grouped_weight else weights_for_gemm + if not inp.requires_grad: + weights_to_save = [None] * len(weights_to_save) - weights_to_save = weights_for_gemm if inp.requires_grad else [None] * num_gemms tensors_to_save, tensor_objects = prepare_for_saving( - grouped_x, + input_to_save, *weights_to_save, split_sizes, base_split_offsets, + input_tensor_offsets, + output_tensor_offsets, ) ctx.save_for_backward(*tensors_to_save) ctx.tensor_objects = tensor_objects - ctx.use_grouped_tensor_path = True + ctx.grouped_tensor_supported = True ctx.weight_quantizers = weight_quantizers ctx.weights_shape_0 = out_features ctx.weights_shape_1 = in_features @@ -418,16 +942,18 @@ def _forward_grouped_tensor( ctx.grad_output_quantizers = grad_output_quantizers ctx.grad_weight_quantizers = grad_weight_quantizers ctx.weights_requires_grad = weight_requires_grad + ctx.single_grouped_weight = single_grouped_weight + ctx.single_grouped_bias = single_grouped_bias if fuse_wgrad_accumulation and ctx.weights_requires_grad: ctx.origin_weight_refs = [weakref.ref(w) for w in weights] ctx.origin_weights_overwrite_main_grad = getattr( weights[0], "overwrite_main_grad", False ) if hasattr(weights[0], "__fsdp_param__"): - ctx.main_grad_funcs = [weights[i].get_main_grad for i in range(num_gemms)] + ctx.main_grad_funcs = [weight.get_main_grad for weight in weights] else: ctx.main_grad_funcs = [ - lambda j=i: weights[j].main_grad for i in range(num_gemms) + lambda j=i: weights[j].main_grad for i in range(len(weights)) ] ctx.device = device ctx.dgrad_out = dgrad_out @@ -451,7 +977,7 @@ def _forward_grouped_tensor( ) ctx.wgrad_store = wgrad_store ctx.debug = False - ctx.save_original_input = False + ctx.save_original_input = save_original_input ctx.input_quantizers = input_quantizers return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces @@ -492,15 +1018,18 @@ def forward( cache_weight, skip_fp8_weight_update, save_original_input, + delayed_scaling_input_quantizer, + unsafe_requantization_input_quantizer, debug, m_splits_tensor, actual_m_splits, unpad_output, + single_grouped_weight, + single_grouped_bias, + use_grouped_tensor, ) = non_tensor_args - if fp8: - backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override - else: - backward_override = None + recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None + backward_override = recipe.backward_override if recipe is not None else None if backward_override == "high_precision": save_original_input = True elif backward_override == "dequantized": @@ -510,25 +1039,46 @@ def forward( use_grouped_gemm_triton = IS_HIP_EXTENSION and os.getenv("NVTE_USE_GROUPED_GEMM_TRITON", "0") == "1" and not fp8 and not fuse_wgrad_accumulation num_gemms = len(m_splits) - weights = weights_and_biases[:num_gemms] - biases = weights_and_biases[num_gemms:] + num_weight_args = 1 if single_grouped_weight else num_gemms + num_bias_args = 1 if single_grouped_bias else num_gemms + weights = weights_and_biases[:num_weight_args] + biases = weights_and_biases[num_weight_args : num_weight_args + num_bias_args] device = inp.device weight_requires_grad = weights[0].requires_grad - # Configure quantizers - if save_original_input and isinstance(input_quantizers[0], Float8Quantizer): - if FP8GlobalStateManager.get_fp8_recipe().custom(): - # Custom recipe factory may produce DS quantizers unknown to caller. - # TODO(negvet): fix on Megatron side — guard should also exclude 'custom', or - # better: check at runtime whether quantizers are DS-based. + origin_weights = weights + is_dist_weight = is_distributed_weight(weights[0]) + if is_dist_weight: + weights = materialize_weight_for_forward(weights) + + backward_needs_input = is_grad_enabled and weight_requires_grad + if backward_override is None and save_original_input and backward_needs_input: + if delayed_scaling_input_quantizer is not None: + if FP8GlobalStateManager.get_fp8_recipe().custom(): + warnings.warn( + "save_original_input is incompatible with delayed-scaling quantizers " + "(Float8Quantizer). Disabling save_original_input for this module.", + stacklevel=2, + ) + save_original_input = False + else: + raise ValueError( + "DelayedScaling recipe is not supported with save_original_input" + ) + + # Megatron-Core may enable this automatically to reuse an activation + # already retained by an upstream operation. The resolved quantizer + # generation is classified once in ``_validate_quantizer_generation``. + if save_original_input and unsafe_requantization_input_quantizer is not None: warnings.warn( - "save_original_input is incompatible with delayed-scaling quantizers " - "(Float8Quantizer). Disabling save_original_input for this module.", + "Ignoring save_original_input=True because the input quantizer cannot " + "safely reconstruct the backward operand from the original input " + f"({unsafe_requantization_input_quantizer}).", stacklevel=2, ) save_original_input = False - else: - raise ValueError("DelayedScaling recipe is not supported with save_original_input") + + # Configure quantizers if input_quantizers[0] is not None: for input_quantizer in input_quantizers: input_quantizer.set_usage( @@ -570,17 +1120,46 @@ def forward( f"weight tensor (shape={tuple(weights[0].size())})" ) - if _GroupedLinear._is_grouped_tensor_path_supported( - fp8=fp8, - fp8_calibration=fp8_calibration, - debug=debug, - cpu_offloading=cpu_offloading, - backward_override=backward_override, - save_original_input=save_original_input, - activation_dtype=activation_dtype, - input_quantizers=input_quantizers, - output_quantizers=output_quantizers, + grouped_tensor_supported = False + if use_grouped_tensor and not ( + fp8_calibration + or debug + or cpu_offloading + or any(q is not None for q in output_quantizers) + ): + if ( + fp8 + and recipe.float8_block_scaling() + and (10, 0) <= get_device_compute_capability() <= (11, 0) + ): + raise RuntimeError( + "use_grouped_tensor=True does not support the FP8 block-scaling recipe on " + "Blackwell GPUs: the native grouped FP8 block-scaling path is Hopper-only. " + "Set use_grouped_tensor=False, or unset " + "NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM if it enabled this path, to use " + "the MXFP8-emulated path on Blackwell." + ) + grouped_tensor_supported = is_module_grouped_tensor_path_supported( + recipe, + activation_dtype, + ) + if ( + use_grouped_tensor + and not grouped_tensor_supported + and (single_grouped_weight or single_grouped_bias) ): + raise RuntimeError( + "Single grouped parameters require the native grouped-tensor path, but the active " + "device, cuBLASLt version, quantization recipe, or GroupedLinear feature " + "configuration does not support it. Disable single_grouped_weight and " + "single_grouped_bias to allow the split-quantize fallback." + ) + if grouped_tensor_supported: + if m_splits.device.type != "cuda": + raise ValueError( + "The native grouped_tensor path requires CUDA m_splits. Pass a CUDA int64 " + "tensor, or set use_grouped_tensor=False." + ) return _GroupedLinear._forward_grouped_tensor( ctx, inp=inp, @@ -600,6 +1179,9 @@ def forward( weight_workspaces=weight_workspaces, cache_weight=cache_weight, skip_fp8_weight_update=skip_fp8_weight_update, + save_original_input=save_original_input, + single_grouped_weight=single_grouped_weight, + single_grouped_bias=single_grouped_bias, weights=weights, biases=biases, out=out, @@ -611,25 +1193,39 @@ def forward( inp_view = inp.reshape(-1, in_features) inputmats: list - if fp8 and not debug: + if IS_HIP_EXTENSION: + if fp8 and not debug: + # Disable bulk allocation when CPU offloading is active: offloading skips small + # tensors (like scales), but bulk allocation shares storage across all tensors, + # so if scales can't be offloaded, nothing in the group can be offloaded. + fused_padding_kwargs = {} + if actual_m_splits is not None and IS_HIP_EXTENSION \ + and inp_view.shape[0] == sum(actual_m_splits): + fused_padding_kwargs["valid_split_sections"] = actual_m_splits + inputmats = tex.split_quantize( + inp_view, m_splits, input_quantizers, + disable_bulk_allocation=cpu_offloading, **fused_padding_kwargs) + elif debug: + inputmats = DebugQuantizer.multi_tensor_quantize( + inp_view, input_quantizers, m_splits, activation_dtype + ) + elif use_grouped_gemm_triton: + inputmats = [cast_if_needed(inp_view, activation_dtype)] + else: + inputmats = torch.split(cast_if_needed(inp_view, activation_dtype), m_splits) + else: # Disable bulk allocation when CPU offloading is active: offloading skips small # tensors (like scales), but bulk allocation shares storage across all tensors, # so if scales can't be offloaded, nothing in the group can be offloaded. - fused_padding_kwargs = {} - if actual_m_splits is not None and IS_HIP_EXTENSION \ - and inp_view.shape[0] == sum(actual_m_splits): - fused_padding_kwargs["valid_split_sections"] = actual_m_splits - inputmats = tex.split_quantize( - inp_view, m_splits, input_quantizers, - disable_bulk_allocation=cpu_offloading, **fused_padding_kwargs) - elif debug: - inputmats = DebugQuantizer.multi_tensor_quantize( - inp_view, input_quantizers, m_splits, activation_dtype + inputmats = _split_quantize( + inp_view, + m_splits, + with_quantized_output=fp8 or debug, + quantizers=input_quantizers, + dtype=activation_dtype, + with_debug_quantizers=debug, + disable_bulk_allocation=cpu_offloading, ) - elif use_grouped_gemm_triton: - inputmats = [cast_if_needed(inp_view, activation_dtype)] - else: - inputmats = torch.split(cast_if_needed(inp_view, activation_dtype), m_splits) if cpu_offloading: start_offload(*inputmats) @@ -717,7 +1313,7 @@ def forward( mark_not_offload(*weights_fp8, *weights) if is_grad_enabled: - ctx.use_grouped_tensor_path = False + ctx.grouped_tensor_supported = False ctx.weight_quantizers = weight_quantizers ctx.weights_shape_1 = weights[0].shape[1] @@ -746,6 +1342,10 @@ def forward( if backward_override == "high_precision" and inp.requires_grad else [None] * num_gemms ) + if is_dist_weight: + # GTP: gathered workspace is transient (re-gathered in backward), don't save it. + weights_fp8 = [None] * num_gemms + saved_weights = origin_weights tensors_to_save, tensor_objects = prepare_for_saving( *inputmats, *weights_fp8, @@ -773,6 +1373,8 @@ def forward( if hasattr(weights[0], "__fsdp_param__"): # MCore FSDP creates main_grad lazily before backward ctx.main_grad_funcs = [weights[i].get_main_grad for i in range(num_gemms)] + elif is_dist_weight: + ctx.main_grad_funcs = [origin_weights[i].grad_buffer for i in range(num_gemms)] else: ctx.main_grad_funcs = [ lambda j=i: weights[j].main_grad for i in range(num_gemms) @@ -833,13 +1435,47 @@ def _backward_grouped_tensor( """Backward path paired with ``_forward_grouped_tensor``.""" saved_tensors = restore_from_func_ctx(ctx) N = ctx.num_gemms - grouped_x = saved_tensors[0] - weights = saved_tensors[1 : 1 + N] - split_sizes = saved_tensors[1 + N] - base_split_offsets = saved_tensors[2 + N] + saved_input = saved_tensors[0] + if ctx.single_grouped_weight: + weights_for_gemm = saved_tensors[1] + weight_tensors = [weights_for_gemm] + split_sizes = saved_tensors[2] + base_split_offsets = saved_tensors[3] + input_tensor_offsets = saved_tensors[4] + output_tensor_offsets = saved_tensors[5] + else: + weight_tensors = saved_tensors[1 : 1 + N] + weights_for_gemm = weight_tensors + split_sizes = saved_tensors[1 + N] + base_split_offsets = saved_tensors[2 + N] + input_tensor_offsets = saved_tensors[3 + N] + output_tensor_offsets = saved_tensors[4 + N] + + if ctx.save_original_input: + x = cast_if_needed( + saved_input.reshape(-1, ctx.weights_shape_1), + ctx.activation_dtype, + ) + if ctx.fp8: + input_quantizer = ctx.input_quantizers[0] + input_quantizer.set_usage(rowwise=False, columnwise=True) + input_quantizer.optimize_for_gemm = True + grouped_x = tex.group_quantize(x, input_quantizer, N, split_sizes) + else: + grouped_x = _GroupedLinear._make_grouped_tensor( + x, + num_gemms=N, + split_sizes=split_sizes, + tensor_offsets=input_tensor_offsets, + last_dim=ctx.weights_shape_1, + dtype=ctx.activation_dtype, + ) + else: + grouped_x = saved_input - origin_weights = [None] * N - main_grads = [None] * N + num_weight_args = 1 if ctx.single_grouped_weight else N + origin_weights = [None] * num_weight_args + main_grads = [None] * num_weight_args if ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: origin_weight_refs = ctx.origin_weight_refs ctx.origin_weight_refs = None @@ -862,12 +1498,18 @@ def _backward_grouped_tensor( columnwise=ctx.weights_requires_grad, ) grad_output_quantizer.optimize_for_gemm = True - if ctx.use_bias and isinstance(grad_output_quantizer, MXFP8Quantizer): + # The grouped FP8 block-scaling bgrad kernel computes dbias in the rowwise + # pass, so the fusion needs rowwise output (i.e. dgrad required). + fuse_bgrad = isinstance(grad_output_quantizer, MXFP8Quantizer) or ( + isinstance(grad_output_quantizer, Float8BlockQuantizer) and ctx.requires_dgrad + ) + if ctx.use_bias and fuse_bgrad: grouped_dy, dbias_packed = tex.bgrad_group_quantize( dy_2d, grad_output_quantizer, N, split_sizes, + tensor_offsets=output_tensor_offsets, ) else: grouped_dy = tex.group_quantize( @@ -875,22 +1517,28 @@ def _backward_grouped_tensor( grad_output_quantizer, N, split_sizes, + tensor_offsets=output_tensor_offsets, ) else: grouped_dy = _GroupedLinear._make_grouped_tensor( dy_2d, num_gemms=N, split_sizes=split_sizes, - base_split_offsets=base_split_offsets, + tensor_offsets=output_tensor_offsets, last_dim=ctx.weights_shape_0, dtype=ctx.activation_dtype, ) - grad_biases = [None] * N if ctx.use_bias: if dbias_packed is None: dbias_packed = compute_grouped_dbias(dy_2d, base_split_offsets, N) - grad_biases = [dbias_packed[i].to(dtype=ctx.activation_dtype) for i in range(N)] + if ctx.single_grouped_bias: + grad_bias_args = [dbias_packed.to(dtype=ctx.activation_dtype)] + else: + grad_bias_args = [dbias_packed[i].to(dtype=ctx.activation_dtype) for i in range(N)] + else: + num_bias_args = 1 if ctx.single_grouped_bias else N + grad_bias_args = [None] * num_bias_args dgrad = None if ctx.requires_dgrad: @@ -899,7 +1547,7 @@ def _backward_grouped_tensor( recipe = ctx.fp8_recipe if hasattr(recipe, "fp8_gemm_dgrad"): dgrad_gemm_use_split_accumulator = recipe.fp8_gemm_dgrad.use_split_accumulator - for weight in weights: + for weight in weight_tensors: if isinstance(weight, QuantizedTensorStorage): weight.update_usage(columnwise_usage=True) dgrad = _GroupedLinear._validate_or_alloc_output( @@ -913,12 +1561,12 @@ def _backward_grouped_tensor( dgrad, num_gemms=N, split_sizes=split_sizes, - base_split_offsets=base_split_offsets, + tensor_offsets=input_tensor_offsets, last_dim=ctx.weights_shape_1, dtype=ctx.activation_dtype, ) general_grouped_gemm_for_grouped_tensor( - weights, + weights_for_gemm, grouped_dy, grouped_dgrad, layout="NN", @@ -939,16 +1587,42 @@ def _backward_grouped_tensor( if hasattr(recipe, "fp8_gemm_wgrad"): wgrad_gemm_use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator if ctx.fuse_wgrad_accumulation: - wgrad_list = main_grads + if ctx.single_grouped_weight: + main_grad = main_grads[0] + grouped_wgrad = GroupedTensor.make_grouped_tensor_from_rowwise_data( + num_tensors=N, + tensor_shape=(ctx.weights_shape_0, ctx.weights_shape_1), + rowwise_data=main_grad.view(-1), + dtype=main_grad.dtype, + ) + wgrad_output = grouped_wgrad + wgrad_list = [main_grad] + else: + wgrad_output = main_grads + wgrad_list = main_grads else: - wgrad_packed = torch.empty( - N, - ctx.weights_shape_0, - ctx.weights_shape_1, - dtype=ctx.activation_dtype, - device=ctx.device, - ) - wgrad_list = [wgrad_packed[i] for i in range(N)] + if ctx.single_grouped_weight: + grouped_wgrad = GroupedTensor.make_grouped_tensor_with_shapes( + num_tensors=N, + shapes=[(ctx.weights_shape_0, ctx.weights_shape_1)] * N, + quantizer=None, + device=ctx.device, + dtype=ctx.activation_dtype, + ) + wgrad_output = grouped_wgrad + wgrad_list = [ + grouped_wgrad.rowwise_data.view(N, ctx.weights_shape_0, ctx.weights_shape_1) + ] + else: + wgrad_packed = torch.empty( + N, + ctx.weights_shape_0, + ctx.weights_shape_1, + dtype=ctx.activation_dtype, + device=ctx.device, + ) + wgrad_output = [wgrad_packed[i] for i in range(N)] + wgrad_list = wgrad_output accumulate = ( accumulate_wgrad_into_param_main_grad @@ -968,9 +1642,9 @@ def grouped_gemm_wgrad(inputmats, grad_output_mats, grad_weights): return None, [None] * N, None if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute(): - ctx.wgrad_store.put([grouped_x, grouped_dy, wgrad_list], grouped_gemm_wgrad) + ctx.wgrad_store.put([grouped_x, grouped_dy, wgrad_output], grouped_gemm_wgrad) else: - grouped_gemm_wgrad(grouped_x, grouped_dy, wgrad_list) + grouped_gemm_wgrad(grouped_x, grouped_dy, wgrad_output) def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): if ctx.weights_requires_grad: @@ -998,10 +1672,7 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): for weight, main_grad, wgrad in zip(origin_weights, main_grads, wgrad_list) ] else: - wgrad_list = [None] * N - - if not ctx.use_bias: - grad_biases = [None] * N + wgrad_list = [None] * num_weight_args if ctx.reduce_and_update_bwd_fp8_tensors: FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) @@ -1012,7 +1683,7 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): None, # out None, # dgrad_out *wgrad_list, - *grad_biases, + *grad_bias_args, ) @staticmethod @@ -1021,7 +1692,7 @@ def backward( ) -> Tuple[Union[torch.Tensor, None], ...]: # pylint: disable=missing-function-docstring with get_nvtx_range_context("_GroupedLinear_backward"): - if ctx.use_grouped_tensor_path: + if ctx.grouped_tensor_supported: return _GroupedLinear._backward_grouped_tensor(ctx, grad_output) saved_tensors = restore_from_func_ctx(ctx) @@ -1037,7 +1708,12 @@ def backward( # Only needed when fuse_wgrad_accumulation is enabled. origin_weights = [None] * N main_grads = [None] * N - if ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: + is_dist_weight = is_distributed_weight(saved_weights[0]) + if is_dist_weight: + origin_weights = saved_weights + if ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: + main_grads = [main_grad_func() for main_grad_func in ctx.main_grad_funcs] + elif ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: origin_weight_refs = ctx.origin_weight_refs ctx.origin_weight_refs = None origin_weights = [ref() if ref is not None else None for ref in origin_weight_refs] @@ -1051,7 +1727,7 @@ def backward( # Preprocess grad output grad_output_view = grad_output.contiguous().view(-1, grad_output.shape[-1]) - + if IS_HIP_EXTENSION: bwd_fused_kwargs = {} if ctx.output_unpadded and ctx.actual_m_splits is not None: bwd_fused_kwargs["valid_split_sections"] = ctx.actual_m_splits @@ -1102,19 +1778,46 @@ def backward( # wgrad GEMM. if not ctx.use_grouped_gemm_triton: grad_output = torch.split( - cast_if_needed(grad_output_view, ctx.activation_dtype), - ctx.m_splits, - ) + cast_if_needed(grad_output_view, ctx.activation_dtype), + ctx.m_splits, + ) else: grad_output = [cast_if_needed(grad_output_view, ctx.activation_dtype)] + else: + grad_output_reference = ctx.grad_output_quantizers[0] + if ctx.fp8 and isinstance(grad_output_reference, HybridQuantizer): + # Usage is a runtime decision, not part of generation validation. + # Apply it uniformly so dispatch can read the first parent without + # rescanning every expert. + for grad_output_quantizer in ctx.grad_output_quantizers: + grad_output_quantizer.set_usage( + rowwise=ctx.requires_dgrad, + columnwise=ctx.weights_requires_grad, + ) + grad_output, grad_biases = _split_quantize_and_bias( + grad_output_view, + ctx.m_splits, + fp8=ctx.fp8, + debug=ctx.debug, + quantizers=ctx.grad_output_quantizers, + dtype=ctx.activation_dtype, + use_bias=ctx.use_bias, + recipe=ctx.fp8_recipe, + disable_bulk_allocation=ctx.cpu_offloading, + ) - if ctx.is_first_microbatch is not None: + if is_dist_weight: + accumulate_wgrad_into_param_main_grad = False + elif ctx.is_first_microbatch is not None: accumulate_wgrad_into_param_main_grad = ( ctx.fuse_wgrad_accumulation and not ctx.is_first_microbatch ) else: accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation + if is_dist_weight: + weights = materialize_weight_for_backward(origin_weights) + if ctx.requires_dgrad: dgrad_gemm_use_split_accumulator = _2X_ACC_DGRAD if ctx.fp8 or ctx.debug: @@ -1200,7 +1903,22 @@ def backward( if ctx.fuse_wgrad_accumulation: wgrad_list = main_grads else: - if not ctx.use_grouped_gemm_triton: + if IS_HIP_EXTENSION: + if not ctx.use_grouped_gemm_triton: + wgrad_packed = torch.empty( + ctx.num_gemms, + *weights[0].size(), + dtype=ctx.activation_dtype, + device=ctx.device, + ) + wgrad_list = [wgrad_packed[i] for i in range(ctx.num_gemms)] + else: + wgrad_list = torch.empty( + (ctx.num_gemms, weights[0].size(0), weights[0].size(1)), + dtype=ctx.activation_dtype, + device=ctx.device + ) + else: wgrad_packed = torch.empty( ctx.num_gemms, *weights[0].size(), @@ -1208,12 +1926,9 @@ def backward( device=ctx.device, ) wgrad_list = [wgrad_packed[i] for i in range(ctx.num_gemms)] - else: - wgrad_list = torch.empty( - (ctx.num_gemms, weights[0].size(0), weights[0].size(1)), - dtype=ctx.activation_dtype, - device=ctx.device - ) + if is_dist_weight: + # Gathered weights are no longer needed after dgrad GEMM. + del weights if ctx.save_original_input: inp = inputmats[0] @@ -1229,28 +1944,39 @@ def backward( else: input_quantizer.set_usage(rowwise=False, columnwise=True) inputmats: list - if ctx.fp8 and not ctx.debug: - save_fused_kwargs = {} - if ctx.actual_m_splits is not None and IS_HIP_EXTENSION \ - and inp_view.shape[0] == sum(ctx.actual_m_splits): - save_fused_kwargs["valid_split_sections"] = ctx.actual_m_splits - inputmats = tex.split_quantize( - inp_view, ctx.m_splits, ctx.input_quantizers, - **save_fused_kwargs) - elif ctx.debug: - inputmats = DebugQuantizer.multi_tensor_quantize( + if IS_HIP_EXTENSION: + if ctx.fp8 and not ctx.debug: + save_fused_kwargs = {} + if ctx.actual_m_splits is not None and IS_HIP_EXTENSION \ + and inp_view.shape[0] == sum(ctx.actual_m_splits): + save_fused_kwargs["valid_split_sections"] = ctx.actual_m_splits + inputmats = tex.split_quantize( + inp_view, ctx.m_splits, ctx.input_quantizers, + **save_fused_kwargs) + elif ctx.debug: + inputmats = DebugQuantizer.multi_tensor_quantize( + inp_view, + ctx.input_quantizers, + ctx.m_splits, + ctx.activation_dtype, + ) + else: + if not ctx.use_grouped_gemm_triton: + inputmats = torch.split( + cast_if_needed(inp_view, ctx.activation_dtype), ctx.m_splits + ) + else: + inputmats = [cast_if_needed(inp_view, ctx.activation_dtype)] + else: + inputmats = _split_quantize( inp_view, - ctx.input_quantizers, ctx.m_splits, - ctx.activation_dtype, + with_quantized_output=ctx.fp8 or ctx.debug, + quantizers=ctx.input_quantizers, + dtype=ctx.activation_dtype, + with_debug_quantizers=ctx.debug, + disable_bulk_allocation=ctx.cpu_offloading, ) - else: - if not ctx.use_grouped_gemm_triton: - inputmats = torch.split( - cast_if_needed(inp_view, ctx.activation_dtype), ctx.m_splits - ) - else: - inputmats = [cast_if_needed(inp_view, ctx.activation_dtype)] elif ctx.backward_override == "dequantized": inputmats_dequant = [] for inputmat in inputmats: @@ -1280,7 +2006,8 @@ def backward( use_split_accumulator=wgrad_gemm_use_split_accumulator, accumulate=( accumulate_wgrad_into_param_main_grad - if not getattr(ctx, "origin_weights_overwrite_main_grad", False) + if not is_dist_weight + and not getattr(ctx, "origin_weights_overwrite_main_grad", False) else False ), **kwargs, @@ -1323,10 +2050,13 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): wgrad = None return wgrad - wgrad_list = [ - handle_custom_ddp_from_mcore(weight, main_grad, wgrad) - for weight, main_grad, wgrad in zip(origin_weights, main_grads, wgrad_list) - ] + if is_dist_weight: + wgrad_list = finalize_weight_grads(origin_weights, wgrad_list) + else: + wgrad_list = [ + handle_custom_ddp_from_mcore(weight, main_grad, wgrad) + for weight, main_grad, wgrad in zip(origin_weights, main_grads, wgrad_list) + ] else: wgrad_list = [None] * ctx.num_gemms @@ -1390,7 +2120,9 @@ class GroupedLinear(TransformerEngineBaseModule): when set to ``True``, this module will not apply the additive bias itself, but instead return the bias value during the forward pass together with the output of the linear transformation :math:`y = xA^T`. This is useful when - the bias addition can be fused to subsequent operations. + the bias addition can be fused to subsequent operations. A single grouped + bias is returned as its packed ``GroupedTensor`` parameter; discrete biases + are returned as a list of per-GEMM tensors. params_dtype : torch.dtype, default = torch.get_default_dtype() it controls the type used to allocate the initial parameters. Useful when the model is trained with lower precision and the original FP32 parameters @@ -1401,7 +2133,8 @@ class GroupedLinear(TransformerEngineBaseModule): If set to ``True``, always saves the original input tensor rather than the cast tensor. In some scenarios, the input tensor is used by multiple modules, and saving the original input tensor may reduce the memory usage. - Cannot work with FP8 DelayedScaling recipe. + Requires input quantizers that can safely reproduce their results from the + original input. Cannot work with FP8 DelayedScaling recipe. single_grouped_weight : bool, default = False If set to ``True``, grouped weights are stored as a single grouped parameter instead of one parameter per GEMM. @@ -1414,6 +2147,13 @@ class GroupedLinear(TransformerEngineBaseModule): EXPERIMENTAL and subject to change. Gated by the ``NVTE_GROUPED_LINEAR_SINGLE_PARAM`` environment variable: if the env var is not set this argument is forced to ``False`` with a warning. + use_grouped_tensor : bool or None, default = None + Prefer the native GroupedTensor grouped GEMM path. Discrete parameters + fall back to split-quantize when the path is unsupported. Single grouped + parameters require the native path and raise instead of falling back. + The native path requires CUDA ``m_splits``. ``None`` preserves the deprecated + ``NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM`` environment-variable + selection for compatibility. New callers should pass a boolean explicitly. Notes ----- @@ -1447,6 +2187,7 @@ def __init__( single_grouped_weight: bool = False, single_grouped_bias: bool = False, name: Optional[str] = None, + use_grouped_tensor: Optional[bool] = None, ) -> None: super().__init__(name) @@ -1462,11 +2203,45 @@ def __init__( self.ub_overlap_ag = ub_overlap_ag self.ub_name = ub_name self.save_original_input = save_original_input + if use_grouped_tensor is None: + use_grouped_tensor_env = os.getenv("NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM") + if use_grouped_tensor_env is not None: + warnings.warn( + "NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM is deprecated and will be " + "removed in a future release. Pass use_grouped_tensor=True or " + "use_grouped_tensor=False to GroupedLinear instead.", + FutureWarning, + stacklevel=2, + ) + else: + use_grouped_tensor_env = "0" + use_grouped_tensor = bool(int(use_grouped_tensor_env)) + if not isinstance(use_grouped_tensor, bool): + raise TypeError( + f"use_grouped_tensor must be a bool or None, got {type(use_grouped_tensor)}." + ) + if IS_HIP_EXTENSION and use_grouped_tensor: + # ROCm has no cuBLASLt grouped-tensor GEMM path; use the split-quantize path. + warnings.warn( + "use_grouped_tensor=True is not supported on ROCm; falling back to the " + "split-quantize grouped GEMM path.", + stacklevel=2, + ) + use_grouped_tensor = False + self.use_grouped_tensor = use_grouped_tensor single_grouped_weight, single_grouped_bias = resolve_grouped_linear_single_param_flags( single_grouped_weight, single_grouped_bias ) self.single_grouped_weight = single_grouped_weight self.single_grouped_bias = single_grouped_bias + if self.use_bias and self.single_grouped_weight and not self.single_grouped_bias: + warnings.warn( + "GroupedLinear has single_grouped_weight=True and bias=True, but " + "single_grouped_bias=False. This requires packing the per-GEMM biases on every " + "forward; enable single_grouped_bias to keep both parameters grouped.", + UserWarning, + stacklevel=2, + ) if ub_overlap_rs or ub_overlap_ag: raise ValueError("GroupedLinear doesn't support Userbuffer overlap.") self.init_method = init_method @@ -1486,6 +2261,9 @@ def __init__( "fwd": 3, "bwd": 2, } + self._validated_quantizer_generations = {} + self._delayed_scaling_input_quantizer = None + self._unsafe_requantization_input_quantizer = None if tp_group is None: self.tp_size = tp_size @@ -1574,6 +2352,57 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: if recipe.float8_current_scaling(): self._customize_quantizers_float8_current_scaling(fwd, recipe) + self._validate_quantizer_generation(fwd) + + def _validate_quantizer_generation(self, fwd: bool) -> None: + """Validate grouped-kernel invariants once per quantizer generation.""" + # Recipe state replaces this list object only when it constructs a new + # quantizer generation. The O(1) identity guard keeps validation off the + # steady-state forward path. Record a generation only after all of its + # operand roles pass, so a failed recipe transition is retried. + meta_key = "scaling_fwd" if fwd else "scaling_bwd" + generation = self.quantizers.get(meta_key) + if generation is None: + return + if self._validated_quantizer_generations.get(meta_key) is generation: + return + + if fwd: + stride = self._num_fp8_tensors_per_gemm["fwd"] + input_quantizers = tuple( + generation[self._offsets["input"] + i * stride] for i in range(self.num_gemms) + ) + weight_quantizers = tuple( + generation[self._offsets["weight"] + i * stride] for i in range(self.num_gemms) + ) + _validate_grouped_quantizer_list(input_quantizers, operand_name="input") + _validate_grouped_quantizer_list(weight_quantizers, operand_name="weight") + delayed_scaling_input_quantizer = next( + (q for q in input_quantizers if isinstance(q, Float8Quantizer)), + None, + ) + unsafe_requantization_input_quantizer = next( + ( + q + for q in input_quantizers + if q is not None and not can_reconstruct_wgrad_input_from_original(q) + ), + None, + ) + self._delayed_scaling_input_quantizer = delayed_scaling_input_quantizer + self._unsafe_requantization_input_quantizer = unsafe_requantization_input_quantizer + else: + stride = self._num_fp8_tensors_per_gemm["bwd"] + grad_output_quantizers = tuple( + generation[self._offsets["grad_output"] + i * stride] for i in range(self.num_gemms) + ) + _validate_grouped_quantizer_list( + grad_output_quantizers, + operand_name="grad_output", + ) + + self._validated_quantizer_generations[meta_key] = generation + def get_quantizer_roles( self, *, @@ -1610,17 +2439,45 @@ def make_grouped_weights(self, defer_init=False) -> None: return weight_quantizers = self._get_weight_quantizers() + # TODO(#3158): Support Identity/Hybrid single grouped weights. + unsupported_quantizers = tuple( + type(quantizer).__name__ + for quantizer in weight_quantizers + if isinstance(quantizer, (IdentityQuantizer, HybridQuantizer)) + ) + if unsupported_quantizers: + quantizer_names = ", ".join(dict.fromkeys(unsupported_quantizers)) + raise NotImplementedError( + "GroupedLinear(single_grouped_weight=True) does not support " + f"{quantizer_names} weight quantizers yet. Set " + "single_grouped_weight=False or unset " + "NVTE_GROUPED_LINEAR_SINGLE_PARAM. See #3158." + ) + recipe = ( weight_quantizers[0]._get_compatible_recipe() if weight_quantizers and weight_quantizers[0] is not None else None ) - if recipe is not None and (recipe.delayed() or recipe.float8_current_scaling()): + if recipe is not None and recipe.delayed(): self.set_tensor_parallel_attributes(defer_init=defer_init) return weights = [getattr(self, f"weight{i}") for i in range(self.num_gemms)] + # TE preserves the original BF16/FP16 initialization on each quantized + # parameter so distributed optimizers can construct lossless FP32 masters. + # Packing the parameters must transfer those values to the new registered + # grouped parameter; otherwise its master is initialized by dequantizing + # MXFP8 and starts from a different value than the discrete-weight layout. + high_precision_init_vals = [_get_high_precision_init_val(weight) for weight in weights] + if any(value is not None for value in high_precision_init_vals) and not all( + value is not None for value in high_precision_init_vals + ): + raise RuntimeError( + "Grouped weights have inconsistent high-precision initialization state" + ) + # Create the weight storage. grouped_weights = GroupedTensor.make_grouped_tensor_with_shapes( num_tensors=self.num_gemms, @@ -1644,9 +2501,18 @@ def make_grouped_weights(self, defer_init=False) -> None: and (weight_quantizers[0] is None or not weight_quantizers[0].internal) ): raise RuntimeError("Found internal quantizer with `single_grouped_weight=True`.") + grouped_parameter = torch.nn.Parameter(grouped_weights) + if all(value is not None for value in high_precision_init_vals): + _attach_high_precision_init_val( + grouped_parameter, + torch.stack(high_precision_init_vals, dim=0), + ) + for weight in weights: + _clear_high_precision_init_val(weight) + self.register_parameter( "weight", - torch.nn.Parameter(grouped_weights), + grouped_parameter, init_fn=self.init_method, get_rng_state_tracker=self.get_rng_state_tracker, fp8_meta_index=self._offsets["weight"], @@ -1892,16 +2758,13 @@ def forward( is_grad_enabled = torch.is_grad_enabled() num_gemms = self.num_gemms - if FP8GlobalStateManager.fp8_graph_capturing(): - skip_fp8_weight_update = ( - FP8GlobalStateManager.quantization_state.skip_fp8_weight_update_tensor - ) - else: - skip_fp8_weight_update = None - if skip_fp8_weight_update is not None: - is_first_microbatch = False - # Make sure splits are in expected format + if (self.single_grouped_weight or self.single_grouped_bias) and not self.use_grouped_tensor: + raise RuntimeError( + "single_grouped_weight and single_grouped_bias require " + "use_grouped_tensor=True; the split-quantize path only supports discrete " + "parameters." + ) if not isinstance(m_splits, torch.Tensor): # Convert list of ints to tensor for backward compatibility m_splits = torch.tensor(m_splits, dtype=torch.int64, device="cpu") @@ -1930,6 +2793,7 @@ def forward( try: weight_tensors = self._get_weight_tensors() bias_tensors = self._get_bias_tensors() + use_grouped_bias = self.use_bias and self.single_grouped_bias quantizers = self._get_quantizers() if not debug else self._get_debug_quantizers() @@ -1937,6 +2801,13 @@ def forward( if self.no_debug_features_active(list(chain(*quantizers))): debug = False quantizers = self._get_quantizers() + if debug and (self.single_grouped_weight or self.single_grouped_bias): + raise RuntimeError( + "TE debug features do not support single grouped parameters. DebugQuantizer " + "uses the split-quantize path, which only supports discrete parameters. " + "Disable single_grouped_weight and single_grouped_bias, or disable TE debug " + "features for this GroupedLinear." + ) ( input_quantizers, @@ -1946,6 +2817,13 @@ def forward( grad_weight_quantizers, grad_output_quantizers, ) = quantizers + if not debug and weight_quantizers[0] is not None: + # Experts share shape and recipe settings: compute once and broadcast. + optimize_for_gemm = self._enable_weight_preswizzle( + weight_quantizers[0], weight_tensors[0] + ) + for q in weight_quantizers: + q.optimize_for_gemm = optimize_for_gemm if is_grad_enabled: linear_fn = _GroupedLinear.apply @@ -1955,11 +2833,14 @@ def forward( autograd_ctx = [None] cache_weight = is_first_microbatch is not None - weight_workspaces = ( - [self._fp8_workspaces.get(f"weight{i}") for i in range(num_gemms)] - if cache_weight - else [None] * num_gemms - ) + if self.single_grouped_weight: + weight_workspaces = [self._fp8_workspaces.get("weight")] if cache_weight else [None] + else: + weight_workspaces = ( + [self._fp8_workspaces.get(f"weight{i}") for i in range(num_gemms)] + if cache_weight + else [None] * num_gemms + ) non_tensor_args = ( self.apply_bias, @@ -1982,10 +2863,15 @@ def forward( cache_weight, skip_fp8_weight_update, self.save_original_input, + self._delayed_scaling_input_quantizer, + self._unsafe_requantization_input_quantizer, debug, m_splits_tensor, actual_m_splits, unpad_output, + self.single_grouped_weight, + use_grouped_bias, + self.use_grouped_tensor, ) out, new_workspaces = linear_fn( *autograd_ctx, @@ -2003,12 +2889,15 @@ def forward( if ws is not None: if isinstance(ws, torch.Tensor): ws = ws.detach() - self._fp8_workspaces[f"weight{i}"] = ws + key = "weight" if self.single_grouped_weight else f"weight{i}" + self._fp8_workspaces[key] = ws finally: self.end_forward() if self.return_bias: + if use_grouped_bias: + return out, bias_tensors[0] return out, [cast_if_needed(b, self.activation_dtype) for b in bias_tensors] return out @@ -2023,31 +2912,34 @@ def backward_dw(self): return with get_nvtx_range_context("_GroupedLinear_wgrad"): (_, grad_biases_, _), tensor_list = self.wgrad_store.pop() - wgrad_list = tensor_list[2] + wgrad_output = tensor_list[2] weight_params = self._get_weight_tensors() if not self.fuse_wgrad_accumulation: - for i in range(self.num_gemms): - weight_params[i].grad = wgrad_list[i].to(weight_params[i].dtype) + if self.single_grouped_weight: + weight_params[0].grad = wgrad_output.rowwise_data.view( + self.num_gemms, self.out_features, self.in_features + ).to(weight_params[0].dtype) + else: + for i in range(self.num_gemms): + weight_params[i].grad = wgrad_output[i].to(weight_params[i].dtype) has_grad_biases = [ grad_bias is not None and grad_bias.numel() != 0 for grad_bias in grad_biases_ ] if self.use_bias and any(has_grad_biases): - grouped_bias = getattr(self, "bias", None) - if grouped_bias is not None: - if not all(has_grad_biases): - raise RuntimeError("Expected all grouped bias gradients to be present.") - gstack = torch.stack(grad_biases_, dim=0).to(grouped_bias.dtype) - if grouped_bias.grad is None: - grouped_bias.grad = gstack - else: - grouped_bias.grad.add_(gstack) - else: - bias_params = [getattr(self, f"bias{i}") for i in range(self.num_gemms)] - for i in range(self.num_gemms): - if has_grad_biases[i] and bias_params[i].grad is None: - bias_params[i].grad = grad_biases_[i].to(bias_params[i].dtype) + if self.use_grouped_tensor: + raise RuntimeError( + "GroupedLinear(use_grouped_tensor=True) fell back to the split-quantize " + "path, which produced per-expert bias gradients during delayed wgrad. " + "This implicit fallback is unsupported with delay_wgrad_compute=True. " + "Use a configuration supported by the grouped-tensor path, or set " + "use_grouped_tensor=False to select the legacy path explicitly." + ) + bias_params = [getattr(self, f"bias{i}") for i in range(self.num_gemms)] + for i in range(self.num_gemms): + if has_grad_biases[i] and bias_params[i].grad is None: + bias_params[i].grad = grad_biases_[i].to(bias_params[i].dtype) del grad_biases_ - del wgrad_list + del wgrad_output del tensor_list self._trigger_wgrad_accumulation_and_reduce_hooks() @@ -2090,10 +2982,7 @@ def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage """Get the weight tensors of the module.""" grouped_weight = getattr(self, "weight", None) if grouped_weight is not None: - weight_tensors = grouped_weight.quantized_tensors - if weight_tensors is None: - # TODO(ksivaman): Remove this after GEMM integration. - weight_tensors = grouped_weight.split_into_quantized_tensors() + weight_tensors = [grouped_weight] else: weight_tensors = [getattr(self, f"weight{i}") for i in range(self.num_gemms)] if not self.fp8 and any(isinstance(w, QuantizedTensorStorage) for w in weight_tensors): @@ -2108,13 +2997,23 @@ def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage return weight_tensors def _get_bias_tensors(self) -> List[torch.Tensor]: - """Per-GEMM bias tensors (views into grouped storage when ``single_grouped_bias``).""" + """Get bias parameters in their registered grouped or per-GEMM layout. + + A single grouped bias remains one packed GroupedTensor. When return_bias=True, + an upper-level framework such as MCore must apply that packed bias accordingly; + Discrete bias parameters retain the existing list-of-per-GEMM contract. + + Example with 2 experts, 128 output features: + + single grouped bias: + GroupedTensor shape = [2, 128] -> [grouped_bias] + + discrete biases: + bias0 [128] + bias1 [128] -> [bias0, bias1] + """ grouped_bias = getattr(self, "bias", None) if grouped_bias is not None: - parts = grouped_bias.quantized_tensors - if parts is None: - parts = grouped_bias.split_into_quantized_tensors() - return [p.reshape(-1) for p in parts] + return [grouped_bias] return [getattr(self, f"bias{i}") for i in range(self.num_gemms)] def _get_weight_quantizers(self) -> List[Quantizer]: @@ -2127,19 +3026,19 @@ def _get_weight_quantizers(self) -> List[Quantizer]: ] for i in range(self.num_gemms) ] - # Preswizzle the weights during quantization instead of lazily inside every GEMM. - # This wont work when primay weights are in fp8 because of 2 reasons - # 1. optimizer step updates would need to dequantize the weights. But swizzled weights - # currently dont support dequantization. - # 2. For FSDP2, quantized weight all-gather would need to be done in the - # unswizzled layout. for i in range(self.num_gemms): weight_quantizers[i].internal = not self.primary_weights_in_fp8 - if not self.primary_weights_in_fp8: - weight_quantizers[i].optimize_for_gemm = True return weight_quantizers def _get_quantizers(self): + if self.fp8: + # Normally validated while installing recipe metadata. Keep this + # O(1) generation guard so failed transitions cannot reuse stale + # validation state if base metadata takes an early return on retry. + self._validate_quantizer_generation(True) + if torch.is_grad_enabled(): + self._validate_quantizer_generation(False) + weight_quantizers = self._get_weight_quantizers() input_quantizers, output_quantizers = ( [None] * self.num_gemms, diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 16fd510b86..8346611795 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -62,6 +62,12 @@ _fsdp_scatter_tensors, _fsdp_gather_tensors, ) +from ..distributed_weight import ( + is_distributed_weight, + materialize_weight_for_forward, + materialize_weight_for_backward, + finalize_weight_grads, +) from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, GemmParallelModes, dist_group_type from ..jit import no_torch_dynamo from ..graph import is_graph_capturing @@ -69,6 +75,7 @@ apply_normalization, noop_cat, set_quantizer_amax_reduction_group, + set_quantizer_usage_for_wgrad_all_gather, WeightGradStore, ) from ..quantized_tensor import ( @@ -81,6 +88,8 @@ from ...debug.pytorch.debug_state import TEDebugState from ..tensor.float8_tensor import Float8CurrentScalingQuantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer +from ..tensor.hybrid_tensor import HybridQuantizer +from ..tensor.identity_tensor import IdentityQuantizer from ..cpu_offload import ( is_cpu_offload_enabled, start_offload, @@ -251,6 +260,8 @@ def forward( # Avoid quantized norm kernel if norm output will be returned # or if a gather of ln_out must be in high precision. custom = is_custom(input_quantizer) + hybrid = isinstance(input_quantizer, HybridQuantizer) + identity = isinstance(input_quantizer, IdentityQuantizer) with_quantized_norm = ( fp8 and not debug @@ -258,6 +269,8 @@ def forward( and not return_layernorm_output_gathered and backward_override is None and not custom # TODO(negvet): and not FP8GlobalStateManager.get_fp8_recipe().custom() + and not hybrid + and not identity ) # ROCm does not currently support quantized norm for Float8CurrentScalingQuantizer @@ -335,6 +348,11 @@ def forward( # ------------------------------------------------------ # Prepare weight tensor # ------------------------------------------------------ + origin_weight = weight + is_dist_weight = is_distributed_weight(origin_weight) + if is_dist_weight: + weight = materialize_weight_for_forward(weight)[0] + out_features = weight.shape[0] new_weight_workspace = None weightmat = weight is_weight_param_quantized = False @@ -544,10 +562,15 @@ def forward( wt_save = weightmat if is_fsdp2 and weightmat is not weight: wt_save = None + # Distributed weight (e.g. GTP): don't save the gathered quantized workspace; + # backward re-gathers from the saved (sharded) weight and re-quantizes. + if is_dist_weight: + wt_save = None + tensors_to_save, tensor_objects = prepare_for_saving( inputmat, wt_save, - weight, + origin_weight, bias, ln_weight, ln_out_to_save, @@ -574,6 +597,8 @@ def forward( if hasattr(weight, "__fsdp_param__"): # MCore FSDP creates main_grad lazily before backward ctx.main_grad_func = weight.get_main_grad + elif is_dist_weight: + ctx.main_grad_func = origin_weight.grad_buffer else: ctx.main_grad_func = lambda: weight.main_grad ctx.grad_input_quantizer = grad_input_quantizer @@ -672,6 +697,9 @@ def backward( rsigma, ) = restore_from_func_ctx(ctx) + is_dist_weight = is_distributed_weight(saved_weight) + if is_dist_weight: + weight = materialize_weight_for_backward(saved_weight)[0] # Restore from weakref to get original weight python object # (preserves attributes like main_grad, grad_added_to_main_grad, etc.) # Only needed when fuse_wgrad_accumulation is enabled. @@ -689,7 +717,7 @@ def backward( ), "weight was removed while fuse_wgrad_accumulation=True" # Since main_grad can be modified inplace, it should not be a part of saved_tensors main_grad = ctx.main_grad_func() if weight is not None else None - if main_grad is not None: + if main_grad is not None and not is_dist_weight: origin_weight.main_grad = main_grad # Gather intermediate/activation tensors if needed @@ -795,12 +823,7 @@ def backward( quantizer = None if ctx.input_quantizer is not None and ctx.fp8: quantizer = ctx.input_quantizer - if quantizer.supports_only_rowwise_all_gather(): - # If data is in FP8, we compute FP8 transposes manually - quantizer.set_usage(rowwise=True, columnwise=False) - else: - # wgrad GEMM requires input with column-wise usage - quantizer.set_usage(rowwise=False, columnwise=True) + set_quantizer_usage_for_wgrad_all_gather(quantizer) if ctx.ub_bulk_dgrad: ln_out_total, _ = fill_userbuffers_buffer_for_all_gather( ub_obj_dgrad, @@ -961,7 +984,7 @@ def backward( and ctx.ub_obj_gradout.with_cublasmp() ): if ctx.grad_output_quantizer is not None: - ctx.grad_output_quantizer.set_usage(rowwise=True, columnwise=False) + set_quantizer_usage_for_wgrad_all_gather(ctx.grad_output_quantizer) grad_output, _ = gather_along_first_dim( grad_output, ctx.tp_group, @@ -1037,7 +1060,10 @@ def backward( use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator # Figure out whether to output wgrad GEMM directly into main grad - if ctx.is_first_microbatch is not None: + if is_dist_weight: + # Distributed weight (e.g. GTP): accumulation happens downstream in finalize. + accumulate_wgrad_into_param_main_grad = False + elif ctx.is_first_microbatch is not None: accumulate_wgrad_into_param_main_grad = ( ctx.fuse_wgrad_accumulation and not ctx.is_first_microbatch ) @@ -1109,6 +1135,9 @@ def wgrad_gemm( # Call wgrad GEMM now wgrad, grad_bias_ = wgrad_gemm(ln_out_total, grad_output) + if is_dist_weight: + wgrad = finalize_weight_grads(saved_weight, [wgrad])[0] + # Update grad bias if needed if grad_bias is None: grad_bias = grad_bias_ @@ -1811,6 +1840,10 @@ def forward( grad_weight_quantizer, grad_output_quantizer, ) = quantizers + if weight_quantizer is not None and not debug: + weight_quantizer.optimize_for_gemm = self._enable_weight_preswizzle( + weight_quantizer, weight_tensor + ) if is_grad_enabled: fwd_fn = _LayerNormLinear.apply @@ -2073,13 +2106,4 @@ def _get_weight_quantizers(self) -> List[Quantizer]: weight_quantizer.set_usage(columnwise=False) else: weight_quantizer.set_usage(columnwise=True if is_nvfp4 else self.keep_fp8_weight_transpose_cache) - else: - # Preswizzle the weights during quantization instead of lazily inside every GEMM. - # This wont work when primay weights are in fp8 because of 2 reasons - # 1. optimizer step updates would need to dequantize the weights. But swizzled weights - # currently dont support dequantization. - # 2. For FSDP2, quantized weight all-gather would need to be done in the - # unswizzled layout. - if not self.primary_weights_in_fp8: - weight_quantizer.optimize_for_gemm = True return [weight_quantizer] diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index ced0066711..f46337348e 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -71,17 +71,20 @@ from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, dist_group_type from ..jit import no_torch_dynamo from ..graph import is_graph_capturing -from ..tensor.float8_tensor import ( - Float8CurrentScalingQuantizer, - Float8Quantizer, - Float8Tensor, -) +from ..tensor.float8_tensor import Float8Tensor from ..tensor.mxfp8_tensor import MXFP8Quantizer if IS_HIP_EXTENSION: from ..tensor.mxfp4_tensor import MXFP4Quantizer from ..tensor.nvfp4_tensor import NVFP4Quantizer from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer -from ._common import apply_normalization, set_quantizer_amax_reduction_group, WeightGradStore +from ..tensor.hybrid_tensor import HybridQuantizer +from ..tensor.identity_tensor import IdentityQuantizer +from ._common import ( + apply_normalization, + set_quantizer_amax_reduction_group, + set_quantizer_usage_for_wgrad_all_gather, + WeightGradStore, +) from ..cpu_offload import ( is_cpu_offload_enabled, start_offload, @@ -439,12 +442,16 @@ def _forward( # for debug: : layernorm output = High precision to enable processing of this norm custom = is_custom(fc1_input_quantizer) + hybrid = isinstance(fc1_input_quantizer, HybridQuantizer) + identity = isinstance(fc1_input_quantizer, IdentityQuantizer) with_quantized_norm = ( fp8 and not debug and not return_layernorm_output and not return_layernorm_output_gathered and not custom + and not hybrid + and not identity ) # ROCm does not currently support quantized norm for Float8CurrentScalingQuantizer @@ -1229,12 +1236,7 @@ def backward( quantizer = None if ctx.fp8 or ctx.debug: quantizer = ctx.fc1_input_quantizer - if isinstance(quantizer, (Float8Quantizer, Float8CurrentScalingQuantizer)): - # If data is in FP8, we compute FP8 transposes manually - quantizer.set_usage(rowwise=True, columnwise=False) - else: - # wgrad GEMM requires input with column-wise usage - quantizer.set_usage(rowwise=False, columnwise=True) + set_quantizer_usage_for_wgrad_all_gather(quantizer) if ctx.ub_bulk_dgrad: ub_obj_fc1_dgrad = get_ub("fc1_dgrad", ctx.fp8) ln_out_total, _ = fill_userbuffers_buffer_for_all_gather( @@ -1355,7 +1357,7 @@ def backward( and ctx.ub_obj_gradout.with_cublasmp() ): if ctx.fc2_grad_output_quantizer is not None: - ctx.fc2_grad_output_quantizer.set_usage(rowwise=True, columnwise=False) + set_quantizer_usage_for_wgrad_all_gather(ctx.fc2_grad_output_quantizer) grad_output, _ = gather_along_first_dim( grad_output, ctx.tp_group, @@ -1532,7 +1534,10 @@ def fc2_wgrad_gemm( if ctx.fp8: # TODO float8 blockwise current scaling (as well as custom quantizers) has no bgrad fusion for now if ( - isinstance(ctx.fc1_grad_output_quantizer, Float8BlockQuantizer) + isinstance( + ctx.fc1_grad_output_quantizer, + (Float8BlockQuantizer, IdentityQuantizer), + ) or ctx.fp8_recipe.custom() ): fc1_bias_grad = dact.view(-1, dact.shape[-1]).sum(dim=0) @@ -2437,6 +2442,15 @@ def forward( fc1_weight, fc2_weight = self._get_weight_tensors() fc1_bias = self.fc1_bias if self.use_bias else None fc2_bias = self.fc2_bias if self.use_bias else None + if not debug: + if fc1_weight_quantizer is not None: + fc1_weight_quantizer.optimize_for_gemm = self._enable_weight_preswizzle( + fc1_weight_quantizer, fc1_weight + ) + if fc2_weight_quantizer is not None: + fc2_weight_quantizer.optimize_for_gemm = self._enable_weight_preswizzle( + fc2_weight_quantizer, fc2_weight + ) if not self.fp8: if isinstance(fc1_weight, Float8Tensor): fc1_weight = fc1_weight.dequantize() @@ -2585,7 +2599,7 @@ def _get_quantizers(self, fp8_output, is_grad_enabled): rowwise=True, columnwise=isinstance( fc2_input_quantizer, - (MXFP8Quantizer, Float8BlockQuantizer, NVFP4Quantizer), + (MXFP8Quantizer, Float8BlockQuantizer, NVFP4Quantizer, HybridQuantizer), ), ) if IS_HIP_EXTENSION and isinstance(fc2_input_quantizer, MXFP4Quantizer): @@ -2822,18 +2836,6 @@ def _get_weight_quantizers(self) -> List[Quantizer]: fc2_weight_quantizer.internal = True if IS_HIP_EXTENSION: fc2_weight_quantizer.set_usage(columnwise = self.keep_fp8_weight_transpose_cache) - else: - # Weight scale factors must be GEMM-swizzled before cuBLAS/CUTLASS can - # consume them. Pre-swizzle once at quantize time (persisted on the cached - # workspace when the weight is cached) instead of lazily inside every GEMM. - # No-op for recipes whose scales don't need swizzling (e.g. per-tensor FP8). - # The exception is quantized_model_init, where the weight parameter is - # itself quantized and gets all-gathered (FSDP2) and optimizer-updated in - # its unswizzled layout; swizzling would break the all-gather and the - # dequantize-on-update, so leave the quantizer state untouched. - if not self.primary_weights_in_fp8: - fc1_weight_quantizer.optimize_for_gemm = True - fc2_weight_quantizer.optimize_for_gemm = True return [fc1_weight_quantizer, fc2_weight_quantizer] def backward_dw(self): diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 618f142f8b..b8a90afb15 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -34,7 +34,13 @@ _2X_ACC_DGRAD, _2X_ACC_WGRAD, ) -from ._common import noop_cat, set_quantizer_amax_reduction_group, WeightGradStore +from ._common import ( + can_reconstruct_wgrad_input_from_original, + noop_cat, + set_quantizer_amax_reduction_group, + set_quantizer_usage_for_wgrad_all_gather, + WeightGradStore, +) from ..quantization import FP8GlobalStateManager, QuantizerRole from ..utils import ( cast_if_needed, @@ -59,6 +65,12 @@ _fsdp_scatter_tensors, _fsdp_gather_tensors, ) +from ..distributed_weight import ( + is_distributed_weight, + materialize_weight_for_forward, + materialize_weight_for_backward, + finalize_weight_grads, +) from ..cpp_extensions import ( general_gemm, ) @@ -281,6 +293,7 @@ def _linear_forward_impl( """ weight = args.weight + is_dist_weight = is_distributed_weight(args.weight) inp = args.inp bias = args.bias input_quantizer = args.input_quantizer @@ -305,10 +318,38 @@ def _linear_forward_impl( is_fsdp2 = args.is_fsdp2 keep_fp8_weight_transpose_cache = args.keep_fp8_weight_transpose_cache use_fsdp2 = args.use_fsdp2 + backward_needs_input = is_grad_enabled and weight.requires_grad if backward_override == "high_precision": save_original_input = True elif backward_override == "dequantized": save_original_input = False + if ( + backward_override is None + and save_original_input + and backward_needs_input + and input_quantizer is not None + ): + # Megatron-Core enables this automatically for attention output projections + # to reuse the high-precision DPA output saved by attention backward. Validate + # the resolved quantizer since this is not necessarily an informed user opt-in. + if isinstance(input_quantizer, Float8Quantizer): + if FP8GlobalStateManager.get_fp8_recipe().custom(): + warnings.warn( + "save_original_input is incompatible with delayed-scaling quantizers " + "(Float8Quantizer). Disabling save_original_input for this module.", + stacklevel=2, + ) + save_original_input = False + else: + raise ValueError("DelayedScaling recipe is not supported with save_original_input") + elif not can_reconstruct_wgrad_input_from_original(input_quantizer): + warnings.warn( + "Ignoring save_original_input=True because the input quantizer cannot " + "safely reconstruct the backward operand from the original input " + f"({input_quantizer}).", + stacklevel=2, + ) + save_original_input = False # NVTX label for profiling nvtx_label = "transformer_engine._Linear.forward" @@ -321,7 +362,6 @@ def _linear_forward_impl( # Configure tensor-parallel communication tp_world_size = get_distributed_world_size(tp_group) - backward_needs_input = is_grad_enabled and weight.requires_grad with_input_all_gather_nccl = ( parallel_mode == "column" and sequence_parallel and not ub_overlap_ag_fprop ) @@ -352,10 +392,6 @@ def _linear_forward_impl( own_quantized_input = False if fp8: assert_dim_for_fp8_exec(inputmat, weight) - if save_original_input: - assert not isinstance( - input_quantizer, Float8Quantizer - ), "DelayedScaling recipe is not supported with save_original_input" if with_input_all_gather_nccl or ub_overlap_ag_fprop: # All-gather input tensor @@ -431,6 +467,14 @@ def _linear_forward_impl( # ------------------------------------------------------ # Prepare weight tensor # ------------------------------------------------------ + # Distributed weight (e.g. GTP): rebind `weight` to the all-gathered tensor; + # `args.weight` keeps the sharded-param reference for backward re-gather / grad + # finalize. No-op for a plain weight. + if is_dist_weight: + weight = materialize_weight_for_forward(args.weight)[0] + # Refresh out_features from the gathered weight (captured sharded above, pre-gather). + out_features = weight.shape[0] + new_weight_workspace = None weightmat = weight if fp8 or debug: @@ -640,6 +684,9 @@ def _linear_forward_impl( wt_save = weightmat if is_fsdp2 and weightmat is not weight: wt_save = None + # Distributed weight (e.g. GTP): don't save the workspace; backward re-gathers it. + if is_dist_weight: + wt_save = None # Dedup save slots that alias forward inputs; ``_linear_setup_ctx`` # rebuilds the refs from ``inp`` / ``weight`` / ``bias``. @@ -749,6 +796,8 @@ def _linear_setup_ctx( bwd_args.origin_weight_overwrites_main_grad = getattr(weight, "overwrite_main_grad", False) if hasattr(weight, "__fsdp_param__"): bwd_args.main_grad_func = weight.get_main_grad + elif is_distributed_weight(weight): + bwd_args.main_grad_func = weight.grad_buffer else: bwd_args.main_grad_func = lambda: weight.main_grad @@ -794,6 +843,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. inputmat = args.inputmat weight_fp8 = args.weight_fp8 saved_weight = args.saved_weight + is_dist_weight = is_distributed_weight(saved_weight) bias = args.bias input_quantizer = args.input_quantizer weight_quantizer = args.weight_quantizer @@ -838,7 +888,8 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. origin_weight_python_object is not None ), "weight was removed while fuse_wgrad_accumulation=True" main_grad = bwd_args.main_grad_func() - origin_weight_python_object.main_grad = main_grad + if not is_dist_weight: + origin_weight_python_object.main_grad = main_grad # Gather intermediate/activation tensors if needed # NOTE: weight_fp8 = weight when bwd_args.fp8 == False and torch.disttributed.FSDP already @@ -974,12 +1025,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. quantizer = None if bwd_args.fp8 or bwd_args.debug: quantizer = input_quantizer - if quantizer.supports_only_rowwise_all_gather(): - # If data is in FP8, we compute FP8 transposes manually - quantizer.set_usage(rowwise=True, columnwise=False) - else: - # wgrad GEMM requires input with column-wise usage - quantizer.set_usage(rowwise=False, columnwise=True) + set_quantizer_usage_for_wgrad_all_gather(quantizer) if bwd_args.ub_bulk_dgrad: inputmat_total, _ = fill_userbuffers_buffer_for_all_gather( ub_obj_dgrad, @@ -1008,6 +1054,12 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. dgrad = None dgrad_work = None + + # Distributed weight (e.g. GTP): re-gather the sharded weight; runs even when + # requires_dgrad=False so the prev_w prefetch is issued for the next layer's bwd. + if is_dist_weight: + weight_fp8 = materialize_weight_for_backward(saved_weight)[0] + if bwd_args.requires_dgrad: # FSDP2: Re-create workspace from all-gathered weight when @@ -1022,6 +1074,16 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. elif bwd_args.weight_quantizer is not None: bwd_args.weight_quantizer.set_usage(rowwise=True, columnwise=True) weight_fp8 = bwd_args.weight_quantizer(saved_weight) + elif ( + is_dist_weight + and bwd_args.fp8 + and bwd_args.weight_quantizer is not None + and not isinstance(weight_fp8, QuantizedTensorStorage) + ): + # Distributed weight re-gathered a BF16 weight: quantize with the layer quantizer + # so the dgrad operand isn't cast by the delayed recipe. + bwd_args.weight_quantizer.set_usage(rowwise=True, columnwise=True) + weight_fp8 = bwd_args.weight_quantizer(weight_fp8) # Make sure required data is available if isinstance(grad_output, QuantizedTensorStorage): @@ -1137,7 +1199,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. and bwd_args.ub_obj_gradout.with_cublasmp() ): if grad_output_quantizer is not None: - grad_output_quantizer.set_usage(rowwise=True, columnwise=False) + set_quantizer_usage_for_wgrad_all_gather(grad_output_quantizer) grad_output, _ = gather_along_first_dim( grad_output, bwd_args.tp_group, @@ -1210,7 +1272,10 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. use_split_accumulator = bwd_args.wgrad_use_split_accumulator # Figure out whether to output wgrad GEMM directly into main grad - if bwd_args.is_first_microbatch is not None: + if is_dist_weight: + # Distributed weight (e.g. GTP): accumulation happens downstream in finalize. + accumulate_wgrad_into_param_main_grad = False + elif bwd_args.is_first_microbatch is not None: accumulate_wgrad_into_param_main_grad = ( bwd_args.fuse_wgrad_accumulation and not bwd_args.is_first_microbatch ) @@ -1286,6 +1351,11 @@ def wgrad_gemm( # Call wgrad GEMM now wgrad, grad_bias_ = wgrad_gemm(inputmat_total, grad_output) + # Distributed weight (e.g. GTP): reduce-scatter the freshly computed wgrad + # (async; overlap with the next layer's bwd via the cascade). + if is_dist_weight: + wgrad = finalize_weight_grads(saved_weight, [wgrad])[0] + # Update grad bias if needed if grad_bias is None: grad_bias = grad_bias_ @@ -1334,15 +1404,19 @@ def wgrad_gemm( origin_weight_python_object, "grad_added_to_main_grad" ): origin_weight_python_object.grad_added_to_main_grad = True + # Use the param's local shape (sharded under GTP) so the dummy wgrad + # matches the saved weight shape; main_grad_func() under GTP returns + # an unsharded scratch and would otherwise mismatch. + wgrad_shape = list(origin_weight_python_object.shape) if getattr(origin_weight_python_object, "zero_out_wgrad", False): wgrad = get_dummy_wgrad( - list(main_grad.shape), + wgrad_shape, origin_weight_python_object.dtype, zero=True, ) else: wgrad = get_dummy_wgrad( - list(main_grad.shape), + wgrad_shape, origin_weight_python_object.dtype, ) elif bwd_args.fuse_wgrad_accumulation: @@ -1551,7 +1625,8 @@ class Linear(TransformerEngineBaseModule): If set to ``True``, always saves the original input tensor rather than the cast tensor. In some scenarios, the input tensor is used by multiple modules, and saving the original input tensor may reduce the memory usage. - Cannot work with FP8 DelayedScaling recipe. + Requires an input quantizer that can safely reproduce its result from the + original input. Cannot work with FP8 DelayedScaling recipe. """ def __init__( @@ -1954,6 +2029,10 @@ def forward( grad_weight_quantizer, grad_output_quantizer, ) = quantizers + if weight_quantizer is not None and not debug: + weight_quantizer.optimize_for_gemm = self._enable_weight_preswizzle( + weight_quantizer, weight_tensor + ) if is_grad_enabled: linear_fn = _Linear.apply @@ -2267,13 +2346,4 @@ def _get_weight_quantizers(self) -> List[Quantizer]: weight_quantizer.set_usage(columnwise=False) else: weight_quantizer.set_usage(columnwise=True if is_nvfp4 else self.keep_fp8_weight_transpose_cache) - else: - # Preswizzle the weights during quantization instead of lazily inside every GEMM. - # This wont work when primay weights are in fp8 because of 2 reasons - # 1. optimizer step updates would need to dequantize the weights. But swizzled weights - # currently dont support dequantization. - # 2. For FSDP2, quantized weight all-gather would need to be done in the - # unswizzled layout. - if not self.primary_weights_in_fp8: - weight_quantizer.optimize_for_gemm = True return [weight_quantizer] diff --git a/transformer_engine/pytorch/newton_schulz.py b/transformer_engine/pytorch/newton_schulz.py index 1cbe6ebfbf..c3509de007 100644 --- a/transformer_engine/pytorch/newton_schulz.py +++ b/transformer_engine/pytorch/newton_schulz.py @@ -2,207 +2,27 @@ # # See LICENSE for license information. -"""Distributed Newton-Schulz matrix orthogonalization via cuSolverMp.""" - -from itertools import chain, cycle, islice, repeat -from typing import Iterator, Literal, Optional, Sequence - -import torch -import torch.distributed as dist - -import transformer_engine_torch as tex - - -_COEFFICIENT_SETS = { - # Values are rounded to closest representable in single precision. - "simple": [ - (3.4445, -4.7750, 2.0315), - ], - "quintic": [ - # optimized for a quintic iteration. - # Source: https://leloykun.github.io/ponder/muon-opt-coeffs/#how-do-we-optimize-the-coefficients - # Numbers from: https://github.com/KellerJordan/modded-nanogpt/blob/0674386070ceb4dcd207e1aca747ffcea6c15250/train_gpt_medium.py#L45 - (4.0848, -6.8946, 2.9270), - (3.9505, -6.3029, 2.6377), - (3.7418, -5.5913, 2.3037), - (2.8769, -3.1427, 1.2046), - (2.8366, -3.0525, 1.2012), - ], - "polar_express": [ - # Polar Express iteration from: https://arxiv.org/abs/2505.16932 - # We include PolarExpress' division by 1.01^polynomial_degree (as stated in their Algorithm 1) in the coefficient list. - # This is a safety factor for numerical stability. - (8.2051, -22.9019, 16.4607), - (4.0664, -2.8612, 0.5184), - (3.9096, -2.8234, 0.5250), - (3.2856, -2.4153, 0.4853), - (2.2779, -1.6198, 0.3985), - (1.8726, -1.2307, 0.3585), - (1.8564, -1.2132, 0.3568), - (1.8750, -1.2500, 0.3750), - ], - "cans": [ - # CANS from: http://arxiv.org/abs/2506.10935 - # CANS iteration (Remez + adaptive interval) based coefficients. - # Source (for generating CANS coefficients): https://github.com/GrishKate/accelerating_orthogonalization/blob/main/polynomials.py - (8.4703, -25.1081, 18.6293), - (4.1828, -3.1087, 0.5806), - (3.9619, -2.9541, 0.5630), - (3.2866, -2.4647, 0.5074), - (2.2737, -1.6447, 0.4162), - ], - "aol": [ - # from https://github.com/thib-s/flash-newton-schulz/blob/main/newton_schulz_triton.py#L511 - (4.0098, -7.0585, 2.4635), - (3.4585, -5.5479, 2.5959), - (2.7573, -3.2939, 1.4254), - (2.7215, -3.0494, 1.3169), - ], -} - -NSCoeffT = Literal[_COEFFICIENT_SETS.keys()] - -CoeffIterMode = Literal["cycle", "repeat_last"] -CoeffT = tuple[float, float, float] - - -def get_coefficient_iterator( - steps: int, - coefficient_sets: Sequence[CoeffT], - mode: CoeffIterMode = "cycle", -) -> Iterator[CoeffT]: - """Iterate through coefficient sets with configurable end behavior using itertools. - - Args: - steps: The number of tuples to yield. - coefficient_sets: A sequence of (a, b, c) coefficient tuples. - mode: Iteration mode: - - "cycle": After the last element, restart from the beginning. - - "repeat_last": After the last element, keep yielding the last tuple. - - Yields: - Tuples (a, b, c) from coefficient_sets according to the specified mode. - - Raises: - ValueError: If coefficient_sets is empty. - ValueError: If an invalid mode is provided. - """ - if not coefficient_sets: - raise ValueError("coefficient_sets must be non-empty.") - - base: Iterator[CoeffT] - if mode == "cycle": - base = cycle(coefficient_sets) - elif mode == "repeat_last": - # Chain the original list with an infinite repeat of the last item - base = chain(coefficient_sets, repeat(coefficient_sets[-1])) - else: - raise ValueError(f"Invalid mode: {mode}. Expected 'cycle' or 'repeat_last'.") - - return islice(base, steps) - - -def get_coefficients(steps: int, coefficient_type: NSCoeffT = "quintic") -> list[CoeffT]: - """Return the coefficient schedule for Newton-Schulz. - - Parameter ``coefficient_type`` can be one of the following - - "simple": Default coefficient set. - - "quintic": Quintic iteration with optimized coefficients. - - "polar_express": Polar Express iteration with optimized coefficients. - - "cans": CANS iteration with Remez + adaptive interval coefficients. - - "aol": AOL coefficient set. - """ - if coefficient_type not in _COEFFICIENT_SETS: - raise ValueError("Invalid coefficient type: " + coefficient_type) - iter_mode: CoeffIterMode = ( - "repeat_last" if coefficient_type in ("polar_express", "cans") else "cycle" - ) - coeff_iter = get_coefficient_iterator( - steps, _COEFFICIENT_SETS[coefficient_type], mode=iter_mode - ) - return list(coeff_iter) - - -class CusolverMpCtx: - """cuSolverMp context for Newton-Schulz matrix orthogonalization. - - Context creation is expensive; create once and reuse across multiple - :func:`newton_schulz` calls. Call :meth:`destroy` when done. - """ - - def __init__(self, group: dist.ProcessGroup) -> None: - self.nranks = dist.get_world_size(group) - self._ptr = tex.cusolvermp_ctx_create( - _get_nccl_comm_ptr(group), dist.get_world_size(group), dist.get_rank(group) - ) - - def destroy(self) -> None: - """Destroy the underlying cuSolverMp context.""" - if self._ptr is not None: - tex.cusolvermp_ctx_destroy(self._ptr) - self._ptr = None - - def __del__(self) -> None: - # Called when the context is manually destroyed or during Python teardown - self.destroy() - - -def _get_nccl_comm_ptr(group: dist.ProcessGroup) -> int: - """Extract the raw NCCL communicator pointer from a PyTorch process group.""" - backend = dist.get_backend(group) - if backend != "nccl": - raise RuntimeError(f"Newton-Schulz requires NCCL backend, got '{backend}'") - nccl_backend = group._get_backend(torch.device("cuda")) - return nccl_backend._comm_ptr() - - -def newton_schulz( - x: torch.Tensor, - ctx: CusolverMpCtx, - num_iterations: int = 5, - coefficients: Optional[Sequence[CoeffT]] = None, -) -> None: - """Compute Newton-Schulz matrix orthogonalization in-place on a distributed matrix. - - Parameters - ---------- - x : torch.Tensor - Local part of the distributed matrix (modified in-place). - Must be a 2D CUDA tensor of type float32 or bfloat16. - Columns are distributed across ranks. - ctx : CusolverMpCtx - cuSolverMp context created by :func:`cusolvermp_ctx_create`. - num_iterations : int, optional - Number of Newton-Schulz iterations. Default: 5. - coefficients : sequence of tuple[float, float, float], optional - Polynomial coefficients for the Newton-Schulz iteration. - """ - if coefficients is None: - coefficients = get_coefficients(num_iterations) - if len(coefficients) != num_iterations: - raise ValueError( - f"Unexpected number of coefficients: {len(coefficients)} for" - f" {num_iterations} iterations" - ) - flat_coefficients: list[float] = [] - for i, coeff in enumerate(coefficients): - if len(coeff) != 3: - raise ValueError( - f"Expected coefficient tuple of length 3 at iteration {i}, got {len(coeff)}" - ) - flat_coefficients.extend(coeff) - - if x.dim() != 2: - raise ValueError(f"Expected 2D tensor, got {x.dim()}D") - if x.dtype not in (torch.float32, torch.bfloat16): - raise ValueError(f"Expected float32 or bfloat16 tensor, got {x.dtype}") - if not x.is_contiguous(): - raise ValueError("Input tensor must be contiguous") - if not x.is_cuda: - raise ValueError("Input tensor must be on CUDA device") - - # Global matrix dimensions; columns are distributed across ranks. - m = x.size(0) - n = x.size(1) * ctx.nranks - - tex.newton_schulz(ctx._ptr, m, n, x, num_iterations, flat_coefficients) +"""Backward-compatible imports for Newton-Schulz orthogonalization.""" + +from transformer_engine.pytorch.optimizers.newton_schulz import ( + CoeffIterMode, + CoeffT, + CusolverMpCtx, + NSCoeffT, + get_coefficient_iterator, + get_coefficients, + newton_schulz, + newton_schulz_tp, +) + + +__all__ = [ + "CoeffIterMode", + "CoeffT", + "CusolverMpCtx", + "NSCoeffT", + "get_coefficient_iterator", + "get_coefficients", + "newton_schulz", + "newton_schulz_tp", +] diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index 607346ce30..f39115c4c6 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -13,11 +13,36 @@ from transformer_engine_torch import FP8TensorMeta from ..torch_version import torch_version from ..quantization import FP8GlobalStateManager +from ..quantized_tensor import QuantizedTensorStorage, Quantizer +from ..tensor import ( + Float8BlockQuantizer, + Float8CurrentScalingQuantizer, + Float8Quantizer, + MXFP8Quantizer, + NVFP4Quantizer, +) from ..tensor.float8_tensor import Float8Tensor -from ..quantized_tensor import QuantizedTensorStorage from ..utils import canonicalize_dtype +def get_fused_normalization_quantizer( + quantizer: Optional[Quantizer], +) -> Optional[Quantizer]: + """Return a quantizer supported by fused normalization kernels.""" + if isinstance( + quantizer, + ( + Float8Quantizer, + Float8CurrentScalingQuantizer, + MXFP8Quantizer, + Float8BlockQuantizer, + NVFP4Quantizer, + ), + ): + return quantizer + return None + + def validate_or_alloc_output( buffer: Optional[torch.Tensor], shape: tuple[int, ...] | list[int], diff --git a/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py index 6def36ffc7..caab9c9c8d 100644 --- a/transformer_engine/pytorch/ops/basic/__init__.py +++ b/transformer_engine/pytorch/ops/basic/__init__.py @@ -24,7 +24,7 @@ from .bias import Bias from .constant_scale import ConstantScale from .dropout import Dropout -from .grouped_linear import GroupedLinear +from .grouped_linear import GroupedLinear, is_op_fuser_grouped_tensor_path_supported from .identity import Identity from .l2normalization import L2Normalization from .layer_norm import LayerNorm diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index f4beffe90c..a974d41ef9 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -6,7 +6,7 @@ from __future__ import annotations import abc -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from typing import Any, Optional import torch @@ -348,18 +348,8 @@ def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: return tex.dsrelu(*args, **kwargs) -class ScaledSReLU(BasicOperation): - r"""Squared ReLU with per-row post-scaling. - - If the SReLU output has shape ``(d_1, ..., d_n)``, it is multiplied - with an extra input tensor of shape ``(d_1, ..., d_{n-1})``. - - Parameters - ---------- - activation_recompute_in_mlp : bool, default = ``False`` - Enable fused grouped MLP kernels to recompute activation outputs - during backward when supported instead of saving them. - """ +class _ScaledUnary(BasicOperation, metaclass=abc.ABCMeta): + """Unary activation with per-row scales (fused grouped MLP middle op).""" num_extra_inputs: int = 1 @@ -367,6 +357,25 @@ def __init__(self, *, activation_recompute_in_mlp: bool = False) -> None: super().__init__() self.activation_recompute_in_mlp: bool = activation_recompute_in_mlp + @abc.abstractmethod + def _scaled_unary_forward( + self, + input_: torch.Tensor, + scales: torch.Tensor, + ) -> torch.Tensor: + """Apply the scaled unary activation.""" + + @abc.abstractmethod + def _scaled_unary_backward( + self, + grad_output: torch.Tensor, + input_: torch.Tensor, + scales: torch.Tensor, + *, + compute_scale_grad: bool, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """Apply the scaled unary activation backward pass.""" + def op_forward(self, *args, **kwargs) -> None: raise RuntimeError( f"{self.__class__.__name__} operation has " @@ -392,7 +401,7 @@ def fuser_forward( prev_op_grad_output_quantizer: Optional[Quantizer], # pylint: disable=unused-argument next_op_input_quantizer: Optional[Quantizer], # pylint: disable=unused-argument basic_op_kwargs: list[dict[str, Any]], # pylint: disable=unused-argument - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: if self.activation_recompute_in_mlp: raise RuntimeError( f"{self.__class__.__name__}(activation_recompute_in_mlp=True) requires the " @@ -410,7 +419,7 @@ def fuser_forward( x = maybe_dequantize(input_.contiguous(), dtype) scales = maybe_dequantize(extra_input, dtype) - y = tex.srelu(x, None) * scales.unsqueeze(-1) + y = self._scaled_unary_forward(x, scales) ctx = basic_op_ctxs[0] if ctx.requires_grad: @@ -448,21 +457,57 @@ def fuser_backward( scales = maybe_dequantize(scales, ctx.dtype) grad_output = maybe_dequantize(grad_output.contiguous(), ctx.dtype) - grad_input = None - if ctx.input_requires_grad: - grad_srelu_out = grad_output * scales.unsqueeze(-1) - grad_input = tex.dsrelu(grad_srelu_out, x, None) - - grad_extra_input = None - if ctx.extra_input_requires_grad: - srelu_out = tex.srelu(x, None) - grad_extra_input = torch.linalg.vecdot(srelu_out, grad_output) + grad_input, grad_extra_input = self._scaled_unary_backward( + grad_output, + x, + scales, + compute_scale_grad=ctx.extra_input_requires_grad, + ) + if not ctx.input_requires_grad: + grad_input = None clear_tensor_data(ctx.saved_tensors[0]) return grad_input, [()], [(grad_extra_input,)] +class ScaledSReLU(_ScaledUnary): + r"""Squared ReLU with per-row post-scaling. + + If the SReLU output has shape ``(d_1, ..., d_n)``, it is multiplied + with an extra input tensor of shape ``(d_1, ..., d_{n-1})``. + + Parameters + ---------- + activation_recompute_in_mlp : bool, default = ``False`` + Enable fused grouped MLP kernels to recompute activation outputs + during backward when supported instead of saving them. + """ + + def _scaled_unary_forward( + self, + input_: torch.Tensor, + scales: torch.Tensor, + ) -> torch.Tensor: + return tex.scaled_srelu(input_, scales, None) + + def _scaled_unary_backward( + self, + grad_output: torch.Tensor, + input_: torch.Tensor, + scales: torch.Tensor, + *, + compute_scale_grad: bool, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + return tex.scaled_dsrelu( + grad_output, + input_, + scales, + None, + compute_scale_grad, + ) + + class SReGLU(_ActivationOperation): r"""Squared Rectified Gated Linear Unit diff --git a/transformer_engine/pytorch/ops/basic/add_extra_input.py b/transformer_engine/pytorch/ops/basic/add_extra_input.py index fc3ca9cade..9af399f2dc 100644 --- a/transformer_engine/pytorch/ops/basic/add_extra_input.py +++ b/transformer_engine/pytorch/ops/basic/add_extra_input.py @@ -5,7 +5,7 @@ """Fusible operation for adding extra input tensor.""" from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from typing import Any, Optional import torch @@ -67,7 +67,7 @@ def fuser_forward( prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: extra_input = basic_op_extra_inputs[0][0] if self._in_place: extra_input = extra_input.detach() diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 57b3b4cd09..4e51cf9873 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -17,7 +17,8 @@ from torch.utils.cpp_extension import IS_HIP_EXTENSION import transformer_engine_torch as tex -from ...constants import DType +from transformer_engine.common.recipe import Recipe +from ...constants import DType, TE_DType from ...cpp_extensions import general_grouped_gemm, general_grouped_gemm_for_grouped_tensor from ...distributed import CudaRNGStatesTracker from ...module._common import WeightGradStore @@ -27,13 +28,12 @@ _2X_ACC_WGRAD, ) from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload -from ...quantization import FP8GlobalStateManager, QuantizerRole, Recipe +from ...quantization import FP8GlobalStateManager, QuantizerRole from ...quantized_tensor import QuantizedTensorStorage from ...tensor import ( - Float8CurrentScalingQuantizer, + Float8BlockQuantizer, MXFP8Quantizer, MXFP8Tensor, - NVFP4Quantizer, Quantizer, ) from ...utils import ( @@ -55,6 +55,12 @@ view_main_grad_as_grouped_buffer, ) from ..op import BasicOperation, OperationContext +from ...distributed_weight import ( + finalize_weight_grads, + is_distributed_weight, + materialize_weight_for_backward, + materialize_weight_for_forward, +) from ...tensor import GroupedTensor, GroupedTensorStorage from ...triton.grouped_dbias_dscales import ( compute_grouped_dbias, @@ -68,6 +74,67 @@ GRAD_INPUT_BUFFER_KEY = "grad_input" +def is_op_fuser_grouped_tensor_path_supported( + recipe: Optional[Recipe], + dtype: torch.dtype, +) -> bool: + """Whether the op-fuser grouped-tensor path supports this recipe and dtype. + + * The graph-safe path dispatches to ``general_grouped_gemm_for_grouped_tensor``, + which is backed by ``nvte_grouped_gemm_with_discrete_inputA`` in the common + library. + * MXFP8 and NVFP4 are supported on Blackwell GPUs with Compute Capability + (CC) 10.x and 11.0. NVFP4 requires RHT because graph-safe grouped + quantization currently requires it. + * FP8 per-tensor current scaling uses grouped current-scaling quantization + through ``tex.group_quantize`` and cuBLASLt grouped GEMM with per-batch + scalar FP8 scaling. It is supported on Hopper and Blackwell, with + cuBLASLt 13.5+ required on Hopper. + * FP8 block scaling uses the grouped-tensor path only on Hopper with + cuBLASLt 13.6+. On other architectures or older cuBLAS versions it + falls back to the split-quantize path for discrete parameters. + * Custom recipes are unsupported because they may assign different + quantizers to input, weight, and grad-output roles. This predicate + currently supports only built-in recipes with known uniform layouts. + * Other quantization recipes, including FP8 delayed scaling, fall back to + split quantization because their grouped quantization kernels are missing. + * Unquantized BF16/FP16 compute is supported on Hopper and Blackwell. FP32 + is excluded because cuBLASLt grouped GEMM does not support it. + * Single grouped parameters have no split-quantize fallback, so callers + must reject them when this function returns ``False``. + """ + if dtype not in (torch.bfloat16, torch.float16): + return False + + device_capability = get_device_compute_capability() + if not (9, 0) <= device_capability <= (11, 0): + return False + cublaslt_version = tex.get_cublasLt_version() + if cublaslt_version < 130300: + return False + if device_capability < (10, 0) and cublaslt_version < 130400: + return False + + if recipe is None: + return True + if recipe.custom(): + return False + if recipe.float8_current_scaling(): + return device_capability >= (10, 0) or cublaslt_version >= 130500 + if recipe.float8_block_scaling(): + # cuBLASLt 13.6 fixes Hopper grouped GEMM algo selection for block-scaled FP8. + return device_capability < (10, 0) and cublaslt_version >= 130600 + if recipe.mxfp8(): + return device_capability >= (10, 0) + if recipe.nvfp4(): + return ( + device_capability >= (10, 0) + and not recipe.disable_rht + and not recipe.row_scaled_activation + ) + return False + + class GroupedLinear(BasicOperation): r"""Apply multiple linear transformations: :math:``y_i = x_i W_i^T + b_i`` @@ -143,11 +210,11 @@ def __init__( delay_wgrad_compute: bool = False, scale_bias: bool = False, ) -> None: - super().__init__() - + # Decide before BasicOperation.__init__ sizes _extra_input_channels. self._scale_bias: bool = scale_bias and bias if self._scale_bias: self.num_extra_inputs = 2 + super().__init__() self.wgrad_store = WeightGradStore(delay_wgrad_compute) self.wgrad_accumulation_and_reduce_hooks: list = [] @@ -223,12 +290,33 @@ def __init__( self._apply_delay_wgrad_param_hooks() + def register_parameter( + self, + name: str, + param: Optional[torch.nn.Parameter], + ) -> None: + """Register a parameter and apply delayed-wgrad metadata when needed.""" + super().register_parameter(name, param) + + # A single-grouped-weight op may be constructed on the meta device and + # receive its grouped parent later. Mark that parent at attachment time, + # before DDP/FSDP inspects the parameter to install backward hooks. + if name == "weight" and param is not None and getattr(self, "single_grouped_weight", False): + wgrad_store = getattr(self, "wgrad_store", None) + if wgrad_store is not None and wgrad_store.delay_wgrad_compute(): + param.skip_backward_post_hook = True + def _apply_delay_wgrad_param_hooks(self) -> None: """Set ``skip_backward_post_hook`` on weights when delaying wgrad (bias uses main backward).""" if not self.wgrad_store.delay_wgrad_compute(): return if self.single_grouped_weight: - self.weight.skip_backward_post_hook = True + # A meta-device op may be created as a parameterless shell and have its + # grouped parent attached after construction. In that case there is no + # ``weight`` parameter to mark yet. + weight = self._parameters.get("weight") + if weight is not None: + weight.skip_backward_post_hook = True else: for group_idx in range(self.num_groups): getattr(self, f"weight{group_idx}").skip_backward_post_hook = True @@ -291,17 +379,28 @@ def backward_dw(self) -> None: w.grad = grad_weights[group_idx].to(w.dtype) self._trigger_wgrad_accumulation_and_reduce_hooks() - def _get_bias_tensors(self, dtype: torch.dtype) -> list[torch.Tensor]: - """Retrieve per-group bias tensors in the given dtype.""" + def _get_discrete_bias_tensors(self, dtype: torch.dtype) -> list[torch.Tensor]: + """Retrieve discrete per-group bias parameters in the given dtype.""" if self.single_grouped_bias: - bias_parts = self.bias.quantized_tensors - if bias_parts is None: - bias_parts = self.bias.split_into_quantized_tensors() - return [maybe_dequantize(p.reshape(-1), dtype) for p in bias_parts] + raise RuntimeError( + "Discrete bias tensors were requested for a single grouped bias parameter." + ) return [ maybe_dequantize(getattr(self, f"bias{idx}"), dtype) for idx in range(self.num_groups) ] + def _get_packed_bias_tensor(self, dtype: torch.dtype) -> torch.Tensor: + """Return all per-group biases as one dense tensor with shape [num_groups, out_features]. + + A single grouped bias is already stored in this layout, so return a view of the + registered parent parameter instead of splitting it into members and stacking it again. + Discrete biases require a stack because they are independent parameters. + """ + if self.single_grouped_bias: + bias_data = self.bias.rowwise_data.view(self.num_groups, self.out_features) + return bias_data if bias_data.dtype == dtype else bias_data.to(dtype=dtype) + return torch.stack(self._get_discrete_bias_tensors(dtype), dim=0) + def num_quantizers(self, mode: str) -> int: if mode == "forward": return 2 * self.num_groups @@ -766,64 +865,6 @@ def op_backward(self, *args, **kwargs): "It overrides `fuser_backward` instead of `op_backward`." ) - @staticmethod - def _is_graph_safe_path_supported( - *, - with_quantized_compute: bool, - input_quantizers: Sequence[Optional[Quantizer]], - dtype: torch.dtype, - single_grouped_weight: bool, - ) -> bool: - """Whether the graph-safe grouped-tensor flow can be used. - - * The graph-safe path dispatches to ``general_grouped_gemm_for_grouped_tensor``, - which is backed by ``nvte_grouped_gemm_with_discrete_inputA`` in the common - library. This filter mirrors cuBLASLt grouped GEMM's architecture - requirement without duplicating its cuBLAS version checks. - * Quantized compute supports MXFP8 and NVFP4 on Blackwell GPUs with Compute Capability (CC) - 10.x and 11.0. NVFP4 requires RHT because graph-safe grouped quantization currently - requires RHT. NVFP4 is additionally restricted to discrete weights: with - ``single_grouped_weight=True`` the weight quantizer is non-RHT and cannot use the - graph-safe grouped quantize kernel, so we fall back to the split-quantize flow. - * FP8 per-tensor current scaling is backed by grouped current-scaling quantization - (``tex.group_quantize``) and cuBLASLt grouped GEMM with per-batch scalar FP8 scaling, - which are supported on Hopper (CC 9.0) and Blackwell (CC 10.x and 11.0). - Every other quantization recipe (fp8 delayed scaling, fp8 block scaling, ...) - falls back to the legacy flow because the corresponding grouped quantization kernels are - missing. - * Unquantized compute supports BF16/FP16 on Hopper (CC 9.0) and Blackwell (CC 10.x and 11.0) - -- FP32 is excluded because the cuBLASLt grouped GEMM doesn't support it. - * Input/weight/grad_output quantizers are assumed to be of the same type, otherwise it - would trigger a fatal error in the cuBLASLt grouped GEMM check. - """ - if IS_HIP_EXTENSION: - # CUDA-only path; ROCm uses general_grouped_gemm. - return False - if not (9, 0) <= get_device_compute_capability() <= (11, 0): - return False - if with_quantized_compute: - # FP8 per-tensor current scaling runs on the Hopper and Blackwell grouped GEMM - # path; the compute-capability range was already checked above. On Hopper it - # requires cuBLAS 13.5+; fall back to the legacy flow on older cuBLAS. - if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers): - if ( - get_device_compute_capability() < (10, 0) - and tex.get_cublasLt_version() < 130500 - ): - return False - return True - # MXFP8 and NVFP4 grouped quantization kernels require Blackwell. - if not (10, 0) <= get_device_compute_capability() <= (11, 0): - return False - if all(isinstance(q, MXFP8Quantizer) for q in input_quantizers): - return True - # NVFP4 graph-safe grouped quantization requires RHT and only supports - # discrete weights; otherwise fall back to the split-quantize flow. - if all(isinstance(q, NVFP4Quantizer) and q.with_rht for q in input_quantizers): - return not single_grouped_weight - return False - return dtype in (torch.bfloat16, torch.float16) - def _get_grouped_weight_for_gemm( self, weight_param: GroupedTensor, @@ -911,6 +952,26 @@ def _get_weight_tensors(self) -> list[torch.nn.Parameter]: return [self.weight] return [getattr(self, f"weight{idx}") for idx in range(self.num_groups)] + def _forward_weight_list(self) -> list[torch.Tensor]: + """Per-expert forward weights, materialized (all-gathered) when distributed.""" + weights = [getattr(self, f"weight{idx}") for idx in range(self.num_groups)] + if is_distributed_weight(weights[0]): + weights = materialize_weight_for_forward(weights) + return weights + + def _backward_weight_setup(self): + """Return ``(origin_weights, is_dist_weight, dgrad_weights)``; dgrad weights are the + re-materialized (all-gathered) weights when distributed, else ``None``.""" + origin_weights = self._get_weight_tensors() + is_dist_weight = is_distributed_weight(origin_weights[0]) + dgrad_weights = materialize_weight_for_backward(origin_weights) if is_dist_weight else None + return origin_weights, is_dist_weight, dgrad_weights + + def _is_distributed_weight(self) -> bool: + """Whether this op's weights are distributed (materialized per fwd/bwd, not saved).""" + leader = self.weight if self.single_grouped_weight else self.weight0 + return is_distributed_weight(leader) + def _get_grouped_bias_for_gemm( self, dtype: torch.dtype, @@ -925,16 +986,7 @@ def _get_grouped_bias_for_gemm( return None num_groups = self.num_groups - if self.single_grouped_bias: - # Already a contiguous (num_groups * out_features) buffer. - bias_data = self.bias.rowwise_data - if bias_data.dtype != dtype: - bias_data = bias_data.to(dtype=dtype) - else: - bias_list = [ - maybe_dequantize(getattr(self, f"bias{idx}"), dtype) for idx in range(num_groups) - ] - bias_data = torch.stack(bias_list, dim=0).contiguous() + bias_data = self._get_packed_bias_tensor(dtype) return GroupedTensorStorage( shape=(num_groups, self.out_features), @@ -954,7 +1006,7 @@ def fuser_forward( prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: num_groups = self.num_groups weight_param = self.weight if self.single_grouped_weight else self.weight0 device = weight_param.device @@ -1003,16 +1055,22 @@ def fuser_forward( out_buffer = basic_op_kwargs[0].get(OUTPUT_BUFFER_KEY) # Dispatch: graph-safe GroupedTensor flow whenever it can be used. - # See ``_is_graph_safe_path_supported`` for the gating rationale -- + # See ``is_op_fuser_grouped_tensor_path_supported`` for the gating rationale -- # in short it requires Hopper (SM90+) plus a supported dtype / # quantization recipe. Otherwise we fall back to the legacy # ``tex.split_quantize`` + ``general_grouped_gemm`` flow. - use_grouped_tensor_path = self._is_graph_safe_path_supported( - with_quantized_compute=with_quantized_compute, - input_quantizers=input_quantizers, - dtype=dtype, - single_grouped_weight=self.single_grouped_weight, + recipe = FP8GlobalStateManager.get_fp8_recipe() if with_quantized_compute else None + use_grouped_tensor_path = is_op_fuser_grouped_tensor_path_supported( + recipe, + dtype, ) + if (self.single_grouped_weight or self.single_grouped_bias) and not use_grouped_tensor_path: + raise RuntimeError( + "Single grouped parameters require the native grouped-tensor GroupedLinear path, " + "which is unavailable for the current device, dtype, or quantization recipe. " + "Disable single_grouped_weight/single_grouped_bias or use a supported grouped-" + "tensor configuration." + ) if use_grouped_tensor_path: out, tensors_to_save = self._fuser_forward_grouped_tensor( @@ -1087,14 +1145,17 @@ def fuser_forward_save_ctx( # temporary workspaces freshly created in each forward pass. if is_cpu_offload_enabled(): saved = tensors_to_save[0] - offset = 4 if self._scale_bias else 3 + # Metadata prefix: + # [split_sizes, base_split_offsets, split_points, + # input_tensor_offsets, output_tensor_offsets, (scales?)] + offset = 6 if self._scale_bias else 5 if use_grouped_tensor_path: - # Layout: [split_sizes, base_split_offsets, split_points, (scales?), grouped_x, *weights] + # Layout: [..., grouped_x, *weights] grouped_x = saved[offset] if grouped_x is not None: mark_activation_offload(grouped_x) else: - # Layout: [split_sizes, None, None, (scales?), *xs, *ws] + # Layout: [..., *xs, *ws] live_xs = [t for t in saved[offset : offset + self.num_groups] if t is not None] if live_xs: mark_activation_offload(*live_xs) @@ -1120,9 +1181,9 @@ def fuser_forward_save_ctx( ctx.weight_quantizers = weight_quantizers ctx.grad_output_quantizers = grad_output_quantizers ctx.grad_input_quantizers = None - # ``split_sizes`` and ``base_split_offsets`` are routed through - # ``save_for_backward`` (see ``_fuser_forward_split_quantize`` and - # ``_fuser_forward_grouped_tensor`` for the saved-tensor layout). + # ``split_sizes``, offset metadata, and related tensors are routed + # through ``save_for_backward`` (see ``_fuser_forward_split_quantize`` + # and ``_fuser_forward_grouped_tensor`` for the saved-tensor layout). if torch.is_autocast_enabled(): ctx.dtype = torch.get_autocast_dtype("cuda") else: @@ -1152,22 +1213,22 @@ def _fuser_forward_split_quantize( out_buffer: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, tuple[Optional[torch.Tensor], ...]]: """Legacy ``tex.split_quantize`` + ``general_grouped_gemm`` flow.""" + if isinstance(input_, GroupedTensor): + raise NotImplementedError( + "Pre-quantized GroupedTensor input is only supported on the " + "graph-safe grouped-tensor path." + ) num_groups = self.num_groups has_bias = self.has_bias # Need CPU split sizes for split_quantize / general_grouped_gemm. split_sizes_int = [int(s) for s in split_sizes.tolist()] - # Extract params - if self.single_grouped_weight: - weights = self.weight.quantized_tensors - if weights is None: - weights = self.weight.split_into_quantized_tensors() - else: - weights = [getattr(self, f"weight{idx}") for idx in range(num_groups)] + # Single grouped parameters are rejected before entering this legacy path. + weights = self._forward_weight_list() # materialized when distributed bs = None if has_bias: - bs = self._get_bias_tensors(dtype) + bs = self._get_discrete_bias_tensors(dtype) ws = self._get_discrete_weights_for_gemm( weights, @@ -1219,7 +1280,8 @@ def _fuser_forward_split_quantize( out_splits[i].add_(bs[i].unsqueeze(0) * scales_splits[i].unsqueeze(-1)) # Prepare weight tensors for backward pass - if not input_requires_grad: + # Distributed weights are re-materialized in backward, so we never save the gathered weight + if not input_requires_grad or self._is_distributed_weight(): ws = [None] * num_groups elif with_quantized_compute: for w, weight_param in zip(ws, weights): @@ -1235,12 +1297,12 @@ def _fuser_forward_split_quantize( # Build the tuple of tensors to save for backward. Layout: # [split_sizes, base_split_offsets, split_points, + # input_tensor_offsets, output_tensor_offsets, # (scales if scale_bias), *xs, *ws] - # ``base_split_offsets`` and ``split_points`` are unused on the - # split-quantize backward path but are included as ``None`` so the - # saved-tensor layout matches the graph-safe - # ``_fuser_forward_grouped_tensor`` path (and the fused MLP forward). - saved: list[Optional[torch.Tensor]] = [split_sizes, None, None] + # Offset metadata slots are unused on the split-quantize backward path + # but are included as ``None`` so the saved-tensor layout matches the + # graph-safe ``_fuser_forward_grouped_tensor`` path. + saved: list[Optional[torch.Tensor]] = [split_sizes, None, None, None, None] if self._scale_bias: saved.append(scales) saved.extend(xs) @@ -1262,28 +1324,61 @@ def _fuser_forward_grouped_tensor( device: torch.device, out_buffer: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, tuple[Optional[torch.Tensor], ...]]: - """Graph-safe GroupedTensor forward path (pure compute). - Returns ``(output, tensors_to_save)``. ``split_sizes``, - ``base_split_offsets`` and ``split_points`` are returned so that - ``fuser_forward_save_ctx`` can call ``save_for_backward`` on them. - """ + """Build graph-safe grouped input storage and run grouped GEMM.""" num_groups = self.num_groups - has_bias = self.has_bias - - base_split_offsets = tex.splits_to_offsets(split_sizes, 1) - split_points = base_split_offsets[1:].to(dtype=torch.int) - - # Flatten to 2D so the first dim is the total token count. + split_sizes, grouped_tensor_offsets = tex.splits_to_offsets_multi( + split_sizes, + device, + strides=[1, 1, self.in_features, self.out_features], + include_leading_zero=[False, True, True, True], + dtypes=[torch.int32, torch.int64, torch.int64, torch.int64], + bulk_allocate=True, + ) + split_points = grouped_tensor_offsets[0] + base_split_offsets = grouped_tensor_offsets[1] + input_tensor_offsets = grouped_tensor_offsets[2] + output_tensor_offsets = grouped_tensor_offsets[3] original_shape = list(input_.size()) - x = maybe_dequantize(input_, dtype).reshape(-1, self.in_features) - total_tokens = x.size(0) - - # Build the input GroupedTensor. + prequantized_input = with_quantized_compute and isinstance(input_, GroupedTensor) if with_quantized_compute: input_quantizer = input_quantizers[0] input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) input_quantizer.optimize_for_gemm = True - grouped_x = tex.group_quantize(x, input_quantizer, num_groups, split_sizes) + if prequantized_input: + # GroupedTensor forbids reshape and is already in the canonical + # (total_tokens, in_features) layout; just validate the shape. + if input_.dim() != 2 or input_.size(-1) != self.in_features: + raise ValueError( + "GroupedTensor input must have shape (total_tokens, " + f"{self.in_features}), but got {tuple(input_.size())}." + ) + total_tokens = input_.size(0) + else: + # Flatten to 2D so the first dim is the total token count. + x = maybe_dequantize(input_, dtype).reshape(-1, self.in_features) + total_tokens = x.size(0) + + # Build the input GroupedTensorStorage for input. + if prequantized_input: + # Input arrived already quantized (e.g. FP8 token dispatch): reuse its rowwise data + # for the GEMM and let the helper supply whatever else the GEMMs need. + grouped_x = input_.copy() + tex.group_requantize_inplace( + grouped_x, + input_quantizer, + num_groups, + split_sizes, + TE_DType[dtype], + tensor_offsets=input_tensor_offsets, + ) + elif with_quantized_compute: + grouped_x = tex.group_quantize( + x, + input_quantizer, + num_groups, + split_sizes, + tensor_offsets=input_tensor_offsets, + ) else: # No quantize: wrap the contiguous high-precision buffer. grouped_x = GroupedTensorStorage( @@ -1293,8 +1388,9 @@ def _fuser_forward_grouped_tensor( quantizer=None, data=x.reshape(-1), first_dims=split_sizes, - tensor_offsets=base_split_offsets * self.in_features, + tensor_offsets=input_tensor_offsets, ) + has_bias = self.has_bias if is_cpu_offload_enabled() and grouped_x is not None: start_offload(grouped_x) @@ -1312,7 +1408,7 @@ def _fuser_forward_grouped_tensor( else: # Discrete weights grouped_weights = self._get_discrete_weights_for_gemm( - [getattr(self, f"weight{idx}") for idx in range(num_groups)], + self._forward_weight_list(), weight_quantizers, columnwise_usage=input_requires_grad, with_quantized_compute=with_quantized_compute, @@ -1329,7 +1425,7 @@ def _fuser_forward_grouped_tensor( quantizer=None, data=out.reshape(-1), first_dims=split_sizes, - tensor_offsets=base_split_offsets * self.out_features, + tensor_offsets=output_tensor_offsets, ) # Bias: hand off to the grouped GEMM (graph-safe, fused). Plain bias @@ -1355,7 +1451,8 @@ def _fuser_forward_grouped_tensor( bias_scale=bias_scale, ) - if not input_requires_grad: + # Distributed weights are re-materialized in backward, so never save the gathered weight. + if not input_requires_grad or self._is_distributed_weight(): grouped_weights = None if self.single_grouped_weight else [None] * num_groups if not weight_requires_grad: @@ -1363,14 +1460,23 @@ def _fuser_forward_grouped_tensor( # Build the tuple of tensors to save for backward. Layout: # [split_sizes, base_split_offsets, split_points, + # input_tensor_offsets, output_tensor_offsets, # (scales if _scale_bias), grouped_x, *weights] + # ``output_tensor_offsets`` matches the linear output row layout and is + # reused as ``grad_output`` offsets in backward. if grouped_x is not None: # (For FP8 per tensor current scaling on Hopper --> Free Rowwise Data # in backward pass) if with_quantized_compute and grouped_x.columnwise_data is not None: grouped_x.rowwise_data = None grouped_x.scale_inv = None - saved: list[Optional[torch.Tensor]] = [split_sizes, base_split_offsets, split_points] + saved: list[Optional[torch.Tensor]] = [ + split_sizes, + base_split_offsets, + split_points, + input_tensor_offsets, + output_tensor_offsets, + ] if self._scale_bias: saved.append(scales) saved.append(grouped_x) @@ -1415,18 +1521,18 @@ def _fuser_backward_split_quantize( ]: num_groups = self.num_groups has_bias = self.has_bias - weights = self._get_weight_tensors() + weights, is_dist_weight, dist_dgrad_weights = self._backward_weight_setup() device = weights[0].device # Saved tensors from forward pass. Layout: # [split_sizes, base_split_offsets, split_points, + # input_tensor_offsets, output_tensor_offsets, # (scales if _scale_bias), *xs, *ws] - # ``base_split_offsets`` and ``split_points`` are unused on this path - # but are present so the saved-tensor layout matches the graph-safe - # path (and the fused MLP forward). + # Offset metadata beyond ``split_sizes`` is unused on this path but is + # present so the saved-tensor layout matches the graph-safe path. saved_tensors = ctx.saved_tensors split_sizes = saved_tensors[0] - saved_tensors = saved_tensors[3:] + saved_tensors = saved_tensors[5:] scales = None if self._scale_bias: scales, saved_tensors = saved_tensors[0], saved_tensors[1:] @@ -1454,7 +1560,7 @@ def _fuser_backward_split_quantize( offsets = torch.zeros(num_groups + 1, dtype=torch.int64, device=device) offsets[1:] = split_sizes.cumsum(0) if self._scale_bias: - bias_packed = torch.stack(self._get_bias_tensors(ctx.dtype)) + bias_packed = self._get_packed_bias_tensor(ctx.dtype) scales_f32 = scales.to(dtype=torch.float32) dbias_packed, grad_scales = compute_grouped_dbias_dscales( dy_2d, @@ -1467,7 +1573,8 @@ def _fuser_backward_split_quantize( grad_biases = [dbias_packed[idx].to(dtype=ctx.dtype) for idx in range(num_groups)] # Initialize grad weight buffers. - accumulate_into_main_grad = self._accumulate_into_main_grad + # Distributed weights reduce their own grads (finalize); never accumulate into main_grad. + accumulate_into_main_grad = self._accumulate_into_main_grad and not is_dist_weight grad_weights = [None] * num_groups final_weight_grads: list[Optional[torch.Tensor]] = ( [None] if self.single_grouped_weight else [None] * num_groups @@ -1515,7 +1622,7 @@ def _fuser_backward_split_quantize( getattr(ctx, "dgrad_out", None), in_shape, ctx.dtype, device ) general_grouped_gemm( - ws, + dist_dgrad_weights if is_dist_weight else ws, dys, [grad_input], [None] * num_groups, # quantization_params @@ -1560,11 +1667,16 @@ def _fuser_backward_split_quantize( if not delay_wgrad: clear_tensor_data(*xs) + # Distributed weights: finalize (e.g. reduce-scatter) the freshly computed wgrads per shard. + # Return discarded (see finalize_weight_grads); the dummy is returned below instead. + if ctx.weight_requires_grad and is_dist_weight: + assert not delay_wgrad, "delayed wgrad unsupported with distributed weights." + finalize_weight_grads(weights, grad_weights) # Megatron-LM wgrad fusion: regardless of overwrite vs. accumulate, # signal that ``main_grad`` already carries the wgrad and replace # ``.grad`` with a dummy so DDP/FSDP hooks won't add ``.grad`` into # ``main_grad`` again. - if ctx.weight_requires_grad and self._accumulate_into_main_grad: + if ctx.weight_requires_grad and (is_dist_weight or self._accumulate_into_main_grad): final_weight_grads = get_dummy_wgrads_for_params(weights) elif ctx.weight_requires_grad and delay_wgrad: final_weight_grads = [None] if self.single_grouped_weight else [None] * num_groups @@ -1596,24 +1708,24 @@ def _fuser_backward_grouped_tensor( Iterable[Iterable[Optional[torch.Tensor]]], Iterable[Iterable[Optional[torch.Tensor]]], ]: + """Graph-safe GroupedTensor backward path.""" num_groups = self.num_groups has_bias = self.has_bias - weights = self._get_weight_tensors() + weights, is_dist_weight, dist_dgrad_weights = self._backward_weight_setup() device = weights[0].device dtype = ctx.dtype - with_quantized_compute = bool(getattr(ctx, "with_quantized_compute", False)) - # Saved tensors from forward pass - # Layout: [split_sizes, base_split_offsets, split_points, - # (scales if _scale_bias), grouped_x, *weights] - # ``split_points`` is unused on this path but is present so the - # saved-tensor layout matches the fused MLP forward (which needs it - # for the cuDNN grouped GEMM kernel). + # Saved tensors from forward pass. Layout: + # [split_sizes, base_split_offsets, split_points, + # input_tensor_offsets, output_tensor_offsets, + # (scales if _scale_bias), grouped_x, *weights] saved_tensors = ctx.saved_tensors split_sizes = saved_tensors[0] base_split_offsets = saved_tensors[1] - saved_tensors = saved_tensors[3:] + input_tensor_offsets = saved_tensors[3] + output_tensor_offsets = saved_tensors[4] + saved_tensors = saved_tensors[5:] scales = None if self._scale_bias: scales, saved_tensors = saved_tensors[0], saved_tensors[1:] @@ -1625,8 +1737,21 @@ def _fuser_backward_grouped_tensor( # Flatten grad_output to 2D (total_tokens, out_features) # to figure out total tokens. - dy_2d = grad_output.reshape(-1, self.out_features) - total_tokens = dy_2d.size(0) + prequantized_grad = with_quantized_compute and isinstance(grad_output, GroupedTensor) + if prequantized_grad: + # GroupedTensor forbids reshape and is already in the canonical + # (total_tokens, out_features) layout; just validate the shape. + if grad_output.dim() != 2 or grad_output.size(-1) != self.out_features: + raise ValueError( + "GroupedTensor grad output must have shape (total_tokens, " + f"{self.out_features}), but got {tuple(grad_output.size())}." + ) + dy_2d = None + total_tokens = grad_output.size(0) + else: + dy_2d = grad_output.reshape(-1, self.out_features) + total_tokens = dy_2d.size(0) + grad_input_shape = list(grad_output.size())[:-1] + [self.in_features] # Build the grad_output GroupedTensor. # Optionally get dbias is fusion available with bgrad_group_quantize @@ -1634,25 +1759,47 @@ def _fuser_backward_grouped_tensor( if with_quantized_compute: grad_output_quantizer = ctx.grad_output_quantizers[0] grad_output_quantizer.set_usage( - rowwise=ctx.input_requires_grad, columnwise=ctx.weight_requires_grad + rowwise=ctx.input_requires_grad, + columnwise=ctx.weight_requires_grad, ) grad_output_quantizer.optimize_for_gemm = True - - if ( - has_bias - and not self._scale_bias - and isinstance(grad_output_quantizer, MXFP8Quantizer) - ): + # FP8 block scaling computes dbias in the rowwise (dgrad) pass, so only fuse + # when dgrad is required. + fuse_bgrad = isinstance(grad_output_quantizer, MXFP8Quantizer) or ( + isinstance(grad_output_quantizer, Float8BlockQuantizer) and ctx.input_requires_grad + ) + if prequantized_grad: + # Grad output arrived already quantized (e.g. FP8 token dispatch): reuse its + # rowwise data for the dgrad GEMM. Bias grads are reduced from the dequantized + # grad below, which is only kept when there is a bias. + grouped_dy = grad_output.copy() + dy_2d = tex.group_requantize_inplace( + grouped_dy, + grad_output_quantizer, + num_groups, + split_sizes, + TE_DType[dtype], + tensor_offsets=output_tensor_offsets, + return_dequantized=has_bias, + ) + elif has_bias and not self._scale_bias and fuse_bgrad: grouped_dy, dbias_packed = tex.bgrad_group_quantize( - dy_2d, grad_output_quantizer, num_groups, split_sizes + dy_2d, + grad_output_quantizer, + num_groups, + split_sizes, + tensor_offsets=output_tensor_offsets, ) else: grouped_dy = tex.group_quantize( - dy_2d, grad_output_quantizer, num_groups, split_sizes + dy_2d, + grad_output_quantizer, + num_groups, + split_sizes, + tensor_offsets=output_tensor_offsets, ) else: dy_2d = maybe_dequantize(dy_2d, dtype) - # Wrap BF16/FP16 buffer as a GroupedTensor for grouped gemm grouped_dy = GroupedTensorStorage( shape=(total_tokens, self.out_features), dtype=dtype, @@ -1660,7 +1807,7 @@ def _fuser_backward_grouped_tensor( quantizer=None, data=dy_2d.reshape(-1), first_dims=split_sizes, - tensor_offsets=base_split_offsets * self.out_features, + tensor_offsets=output_tensor_offsets, ) # Bias Grads compute if not already computed in bgrad_group_quantize @@ -1668,7 +1815,7 @@ def _fuser_backward_grouped_tensor( grad_scales: Optional[torch.Tensor] = None if has_bias: if self._scale_bias: - bias_packed = torch.stack(self._get_bias_tensors(dtype)) + bias_packed = self._get_packed_bias_tensor(dtype) scales_f32 = scales.to(dtype=torch.float32) dbias_packed, grad_scales = compute_grouped_dbias_dscales( dy_2d, @@ -1677,7 +1824,8 @@ def _fuser_backward_grouped_tensor( offsets=base_split_offsets, ) elif dbias_packed is None: - # BF16/FP16 path + # BF16/FP16 and pre-quantized MXFP8 paths, neither of which fuses dbias + # into a quantize kernel. dbias_packed = compute_grouped_dbias(dy_2d, base_split_offsets, num_groups) if self.single_grouped_bias: final_bias_grads = [dbias_packed.to(dtype=dtype)] @@ -1687,7 +1835,6 @@ def _fuser_backward_grouped_tensor( # ---- dgrad GEMM ---------------------------------------------------- grad_input = None if ctx.input_requires_grad: - grad_input_shape = list(grad_output.size())[:-1] + [self.in_features] grad_input = validate_or_alloc_output( getattr(ctx, "dgrad_out", None), grad_input_shape, dtype, device ) @@ -1698,10 +1845,10 @@ def _fuser_backward_grouped_tensor( quantizer=None, data=grad_input.reshape(-1), first_dims=split_sizes, - tensor_offsets=base_split_offsets * self.in_features, + tensor_offsets=input_tensor_offsets, ) general_grouped_gemm_for_grouped_tensor( - ws, + dist_dgrad_weights if is_dist_weight else ws, grouped_dy, grouped_grad_input, layout="NN", @@ -1747,7 +1894,8 @@ def _fuser_backward_grouped_tensor( final_weight_grads[0] = grouped_wgrad.rowwise_data.view(num_groups, *weight_shape) wgrad_output = grouped_wgrad else: - if self._accumulate_into_main_grad: + # Distributed weights finalize wgrads (below); never accumulate into main_grad. + if self._accumulate_into_main_grad and not is_dist_weight: final_weight_grads = [ get_main_grad_from_param(w, op_label="GroupedLinear") for w in weights ] @@ -1777,11 +1925,16 @@ def _fuser_backward_grouped_tensor( else: wgrad_gemm(grouped_x, grouped_dy, wgrad_output) + # Distributed weights: finalize (e.g. reduce-scatter) the freshly computed wgrads per shard. + # Return discarded (see finalize_weight_grads); the dummy is returned below instead. + if ctx.weight_requires_grad and is_dist_weight: + assert not delay_wgrad, "delayed wgrad unsupported with distributed weights." + finalize_weight_grads(weights, final_weight_grads) # Megatron-LM wgrad fusion: regardless of overwrite vs. accumulate, # signal that ``main_grad`` already carries the wgrad and replace # ``.grad`` with a dummy so DDP/FSDP hooks won't add ``.grad`` into # ``main_grad`` again. - if ctx.weight_requires_grad and self._accumulate_into_main_grad: + if ctx.weight_requires_grad and (is_dist_weight or self._accumulate_into_main_grad): final_weight_grads = get_dummy_wgrads_for_params(weights) elif ctx.weight_requires_grad and delay_wgrad: final_weight_grads = [None] if self.single_grouped_weight else [None] * num_groups diff --git a/transformer_engine/pytorch/ops/basic/layer_norm.py b/transformer_engine/pytorch/ops/basic/layer_norm.py index b15dd36604..4f91e6e055 100644 --- a/transformer_engine/pytorch/ops/basic/layer_norm.py +++ b/transformer_engine/pytorch/ops/basic/layer_norm.py @@ -29,7 +29,11 @@ devices_match, ) from ..op import BasicOperation, OperationContext -from .._common import maybe_autocast_dtype, maybe_dequantize +from .._common import ( + get_fused_normalization_quantizer, + maybe_autocast_dtype, + maybe_dequantize, +) class LayerNorm(BasicOperation): @@ -188,6 +192,9 @@ def op_forward( if is_in_onnx_export_mode(): return self.op_onnx_forward(input_) + # Fall back to a high-precision output when fused quantization is unsupported. + next_op_input_quantizer = get_fused_normalization_quantizer(next_op_input_quantizer) + # Check tensor dims weight = self.weight weight_dims = tuple(weight.size()) diff --git a/transformer_engine/pytorch/ops/basic/make_extra_output.py b/transformer_engine/pytorch/ops/basic/make_extra_output.py index 0d9c870262..47b7c6495d 100644 --- a/transformer_engine/pytorch/ops/basic/make_extra_output.py +++ b/transformer_engine/pytorch/ops/basic/make_extra_output.py @@ -5,7 +5,7 @@ """Make extra tensor output in operation fuser.""" from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from typing import Any, Optional import torch @@ -72,7 +72,7 @@ def fuser_forward( prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: return input_, [(input_,)] def fuser_backward( @@ -80,14 +80,17 @@ def fuser_backward( basic_op_ctxs: list[OperationContext], grad_output: torch.Tensor, *, - basic_op_grad_extra_outputs: list[tuple[torch.Tensor, ...]], + basic_op_grad_extra_outputs: list[tuple[Optional[torch.Tensor], ...]], ) -> tuple[ torch.Tensor, Iterable[Iterable[Optional[torch.Tensor]]], Iterable[Iterable[Optional[torch.Tensor]]], ]: grad_extra_output = basic_op_grad_extra_outputs[0][0] - if self._in_place: + if grad_extra_output is None: + # Extra output is not consumed, so its gradient is zero + grad_input = grad_output + elif self._in_place: grad_extra_output += grad_output grad_input = grad_extra_output else: diff --git a/transformer_engine/pytorch/ops/basic/rmsnorm.py b/transformer_engine/pytorch/ops/basic/rmsnorm.py index 0491ab9143..db976e81d1 100644 --- a/transformer_engine/pytorch/ops/basic/rmsnorm.py +++ b/transformer_engine/pytorch/ops/basic/rmsnorm.py @@ -32,7 +32,11 @@ devices_match, ) from ..op import BasicOperation, OperationContext -from .._common import maybe_autocast_dtype, maybe_dequantize +from .._common import ( + get_fused_normalization_quantizer, + maybe_autocast_dtype, + maybe_dequantize, +) class RMSNorm(BasicOperation): @@ -175,6 +179,9 @@ def op_forward( if is_in_onnx_export_mode(): return self.op_onnx_forward(input_) + # Fall back to a high-precision output when fused quantization is unsupported. + next_op_input_quantizer = get_fused_normalization_quantizer(next_op_input_quantizer) + # Check tensor dims weight = self.weight weight_dims = tuple(weight.size()) diff --git a/transformer_engine/pytorch/ops/basic/swiglu.py b/transformer_engine/pytorch/ops/basic/swiglu.py index 02f330ede3..598126e2cb 100644 --- a/transformer_engine/pytorch/ops/basic/swiglu.py +++ b/transformer_engine/pytorch/ops/basic/swiglu.py @@ -5,7 +5,7 @@ """Fusible operation for SwiGLU and variants.""" from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from typing import Any, Optional import torch @@ -387,14 +387,21 @@ def __init__( self.glu_interleave_size: Optional[int] = glu_interleave_size self.activation_recompute_in_mlp: bool = activation_recompute_in_mlp - def _glu_forward(self, swiglu_in: torch.Tensor) -> torch.Tensor: + def _scaled_glu_forward( + self, + input_: torch.Tensor, + scales: torch.Tensor, + ) -> torch.Tensor: raise NotImplementedError - def _glu_backward( + def _scaled_glu_backward( self, - grad_swiglu_out: torch.Tensor, - swiglu_in: torch.Tensor, - ) -> torch.Tensor: + grad_output: torch.Tensor, + input_: torch.Tensor, + scales: torch.Tensor, + *, + compute_scale_grad: bool, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: raise NotImplementedError def op_forward(self, *args, **kwargs) -> None: @@ -422,7 +429,7 @@ def fuser_forward( prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: if self.activation_recompute_in_mlp: raise RuntimeError( f"{self.__class__.__name__}(activation_recompute_in_mlp=True) requires the " @@ -442,22 +449,7 @@ def fuser_forward( # Make sure inputs are in correct dtype input_ = maybe_dequantize(input_, dtype) scales = maybe_dequantize(extra_input, dtype) - - # Remove gate interleaving if needed - swiglu_in = input_ - if self.glu_interleave_size is not None: - shape = swiglu_in.size() - swiglu_in = swiglu_in.reshape( - -1, - shape[-1] // (2 * self.glu_interleave_size), - 2, - self.glu_interleave_size, - ) - swiglu_in = swiglu_in.transpose(1, 2).contiguous() - swiglu_in = swiglu_in.view(shape) - - swiglu_out = self._glu_forward(swiglu_in) - out = swiglu_out * scales.unsqueeze(-1) + out = self._scaled_glu_forward(input_, scales) # Save state for backward pass ctx = basic_op_ctxs[0] @@ -469,7 +461,7 @@ def fuser_forward( ctx.dtype = dtype ctx.save_for_backward( input_, - scales if ctx.input_requires_grad else None, + scales if ctx.input_requires_grad or ctx.extra_input_requires_grad else None, ) return out, [()] @@ -498,41 +490,14 @@ def fuser_backward( scales = maybe_dequantize(scales, ctx.dtype) grad_output = maybe_dequantize(grad_output, ctx.dtype) - # Remove gate interleaving if needed - swiglu_in = input_ - if self.glu_interleave_size is not None: - shape = swiglu_in.size() - swiglu_in = swiglu_in.reshape( - -1, - shape[-1] // (2 * self.glu_interleave_size), - 2, - self.glu_interleave_size, - ) - swiglu_in = swiglu_in.transpose(1, 2).contiguous() - swiglu_in = swiglu_in.view(shape) - - # Compute input grad - grad_input = None - if ctx.input_requires_grad: - grad_swiglu_out = grad_output * scales.unsqueeze(-1) - grad_swiglu_in = self._glu_backward(grad_swiglu_out, swiglu_in) - grad_input = grad_swiglu_in - if self.glu_interleave_size is not None: - shape = grad_input.size() - grad_input = grad_input.reshape( - -1, - 2, - shape[-1] // (2 * self.glu_interleave_size), - self.glu_interleave_size, - ) - grad_input = grad_input.transpose(1, 2).contiguous() - grad_input = grad_input.view(shape) - - # Compute scales grad by recomputing GLU - grad_extra_input = None - if ctx.extra_input_requires_grad: - swiglu_out = self._glu_forward(swiglu_in) - grad_extra_input = torch.linalg.vecdot(swiglu_out, grad_output) + grad_input, grad_extra_input = self._scaled_glu_backward( + grad_output, + input_, + scales, + compute_scale_grad=ctx.extra_input_requires_grad, + ) + if not ctx.input_requires_grad: + grad_input = None # Clear input tensor if possible clear_tensor_data(ctx.saved_tensors[0]) # input_ @@ -558,15 +523,34 @@ class ScaledSwiGLU(_ScaledGLU): """ - def _glu_forward(self, swiglu_in: torch.Tensor) -> torch.Tensor: - return tex.swiglu(swiglu_in, None) - - def _glu_backward( + def _scaled_glu_forward( self, - grad_swiglu_out: torch.Tensor, - swiglu_in: torch.Tensor, + input_: torch.Tensor, + scales: torch.Tensor, ) -> torch.Tensor: - return tex.dswiglu(grad_swiglu_out, swiglu_in, None) + return tex.scaled_swiglu( + input_, + scales, + None, + int(self.glu_interleave_size or 0), + ) + + def _scaled_glu_backward( + self, + grad_output: torch.Tensor, + input_: torch.Tensor, + scales: torch.Tensor, + *, + compute_scale_grad: bool, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + return tex.scaled_dswiglu( + grad_output, + input_, + scales, + None, + int(self.glu_interleave_size or 0), + compute_scale_grad, + ) class ScaledClampedQGeGLU(_ScaledGLU): @@ -614,16 +598,39 @@ def __init__( glu_linear_offset=glu_linear_offset, ) - def _glu_forward(self, swiglu_in: torch.Tensor) -> torch.Tensor: - return self._clamped._tex_clamped_swiglu_forward(swiglu_in, None) - - def _glu_backward( + def _scaled_glu_forward( self, - grad_swiglu_out: torch.Tensor, - swiglu_in: torch.Tensor, + input_: torch.Tensor, + scales: torch.Tensor, ) -> torch.Tensor: - return self._clamped._tex_clamped_dswiglu( - grad_swiglu_out, - swiglu_in, + clamped = self._clamped + return tex.scaled_clamped_swiglu( + input_, + scales, + None, + clamped.limit, + clamped.alpha, + clamped.glu_linear_offset, + int(self.glu_interleave_size or 0), + ) + + def _scaled_glu_backward( + self, + grad_output: torch.Tensor, + input_: torch.Tensor, + scales: torch.Tensor, + *, + compute_scale_grad: bool, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + clamped = self._clamped + return tex.scaled_clamped_dswiglu( + grad_output, + input_, + scales, None, + clamped.limit, + clamped.alpha, + clamped.glu_linear_offset, + int(self.glu_interleave_size or 0), + compute_scale_grad, ) diff --git a/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py b/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py index a3c81e60c8..483749da47 100644 --- a/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py +++ b/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py @@ -33,7 +33,7 @@ def fuser_backward( basic_op_ctxs: list[OperationContext], grad_output: torch.Tensor, *, - basic_op_grad_extra_outputs: list[tuple[torch.Tensor, ...]], + basic_op_grad_extra_outputs: list[tuple[Optional[torch.Tensor], ...]], ) -> tuple[ torch.Tensor, list[tuple[Optional[torch.Tensor], ...]], @@ -56,7 +56,11 @@ def fuser_backward( extra_grad = basic_op_grad_extra_outputs[0][0] dy = maybe_dequantize(grad_output.contiguous(), dtype).view(x.size()) w = maybe_dequantize(rmsnorm_op.weight, dtype).view((inner_dim,)) - add = maybe_dequantize(extra_grad.contiguous(), dtype).view(x.size()) + if extra_grad is None: + # Extra output is not consumed, so its gradient is zero + add = torch.zeros_like(dy) + else: + add = maybe_dequantize(extra_grad.contiguous(), dtype).view(x.size()) # Compute RMSNorm backward pass dx, dw = tex.rmsnorm_bwd_add( diff --git a/transformer_engine/pytorch/ops/fused/backward_linear_add.py b/transformer_engine/pytorch/ops/fused/backward_linear_add.py index 382fecfd07..60b1320b2e 100644 --- a/transformer_engine/pytorch/ops/fused/backward_linear_add.py +++ b/transformer_engine/pytorch/ops/fused/backward_linear_add.py @@ -40,7 +40,7 @@ def fuser_backward( basic_op_ctxs: list[OperationContext], grad_output: torch.Tensor, *, - basic_op_grad_extra_outputs: list[tuple[torch.Tensor, ...]], + basic_op_grad_extra_outputs: list[tuple[Optional[torch.Tensor], ...]], ) -> tuple[ torch.Tensor, list[tuple[Optional[torch.Tensor], ...]], @@ -67,19 +67,21 @@ def fuser_backward( else: accumulate_into_main_grad = False - # Linear backward pass + # Linear backward pass. Skip in-place add when there is no + # extra-output gradient. grad_input = basic_op_grad_extra_outputs[0][0] + accumulate_into_grad_input = grad_input is not None grad_input, grad_weight = BasicLinear._functional_backward( grad_output=grad_output, input=x_local, weight=w, input_requires_grad=linear_op_ctx.input_requires_grad, weight_requires_grad=linear_op_ctx.weight_requires_grad, - dtype=grad_input.dtype, + dtype=None if grad_input is None else grad_input.dtype, grad_weight=grad_weight, accumulate_into_grad_weight=accumulate_into_main_grad, grad_input=grad_input, - accumulate_into_grad_input=True, + accumulate_into_grad_input=accumulate_into_grad_input, tensor_parallel_mode=linear_op.tensor_parallel_mode, tensor_parallel_group=linear_op.tensor_parallel_group, sequence_parallel=linear_op.sequence_parallel, diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py b/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py index 8df929f799..6fa63675b4 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py @@ -5,7 +5,7 @@ """Fused operation for forward GEMM + bias + activation.""" from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Sequence from typing import Any, Optional import torch @@ -59,7 +59,7 @@ def fuser_forward( prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: # Get basic operations idx = self._op_idxs["linear"] diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py index 5376a7d264..28586360f5 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py @@ -5,7 +5,7 @@ """Fused operation for forward GEMM + bias + add.""" from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Sequence from typing import Any, Optional import torch @@ -57,7 +57,7 @@ def fuser_forward( prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: # Get basic operations idx = self._op_idxs["linear"] diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py index abeb39adfa..277263e0ec 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py @@ -5,7 +5,7 @@ """Fused operation for forward GEMM + scale + add.""" from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Sequence from typing import Any, Optional import torch @@ -47,7 +47,7 @@ def fuser_forward( prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: # Get basic operations linear_op = self.basic_ops[0] diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 31189af09c..6b3f53fbd9 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -6,7 +6,7 @@ from __future__ import annotations -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Sequence import functools import os from importlib.metadata import PackageNotFoundError, version as get_pkg_version @@ -16,16 +16,22 @@ from packaging.version import Version as PkgVersion import transformer_engine_torch as tex -from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE +from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor +from ...distributed_weight import ( + is_distributed_weight, + materialize_weight_for_forward, + materialize_weight_for_backward, + finalize_weight_grads, +) from ...module.base import _2X_ACC_WGRAD from ...quantization import Recipe from ...tensor import NVFP4Quantizer, NVFP4Tensor, NVFP4TensorStorage, Quantizer from ...tensor.grouped_tensor import GroupedTensor -from ...tensor.mxfp8_tensor import MXFP8Quantizer +from ...tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor from ...tensor.storage.grouped_tensor_storage import GroupedTensorStorage -from ...triton.grouped_dbias_dscales import compute_grouped_dbias_dscales +from ...triton.grouped_dbias_dscales import compute_grouped_dbias, compute_grouped_dbias_dscales from ...utils import ( ceil_div, clear_tensor_data, @@ -95,15 +101,24 @@ def _nvidia_cudnn_frontend_supports_wgrad() -> bool: return _cudnn_frontend_version_supported() -def _wrap_single_nvfp4_as_grouped( +def _cudnn_frontend_supports_single_group_runtime_offsets( + activation_type: type[FusibleOperation], +) -> bool: + """Check cuDNN FE support for single-group runtime offsets.""" + return not issubclass(activation_type, ScaledSReLU) and _cudnn_frontend_version_at_least( + "1.27.0" + ) + + +def _wrap_single_quantized_as_grouped( tensor: torch.Tensor, - quantized: NVFP4Tensor | NVFP4TensorStorage, - quantizer: NVFP4Quantizer, + quantized: MXFP8Tensor | NVFP4Tensor | NVFP4TensorStorage, + quantizer: MXFP8Quantizer | NVFP4Quantizer, split_sizes: Optional[torch.Tensor], *, tensor_offsets: Optional[torch.Tensor] = None, ) -> GroupedTensor: - """Wrap a single NVFP4 tensor in GroupedTensor storage.""" + """Wrap a single quantized tensor in GroupedTensor storage.""" with_gemm_swizzled_scales = quantized._with_gemm_swizzled_scales if quantizer.optimize_for_gemm: tex.swizzle_scales_for_gemm_(quantized) @@ -113,8 +128,8 @@ def _wrap_single_nvfp4_as_grouped( rowwise_scale = quantized._rowwise_scale_inv columnwise_data = quantized._columnwise_data columnwise_scale = quantized._columnwise_scale_inv - amax = quantized._amax_rowwise - columnwise_amax = quantized._amax_columnwise + amax = getattr(quantized, "_amax_rowwise", None) + columnwise_amax = getattr(quantized, "_amax_columnwise", None) if split_sizes is None: split_sizes = torch.full((1,), tensor.shape[0], dtype=torch.int64, device=tensor.device) @@ -122,21 +137,13 @@ def _wrap_single_nvfp4_as_grouped( split_sizes = split_sizes.to(dtype=torch.int64, device=tensor.device) m_dim = tensor.shape[0] - if rowwise_data is not None: + if isinstance(quantizer, NVFP4Quantizer) and rowwise_data is not None: k_dim = rowwise_data.shape[-1] * 2 - elif columnwise_data is not None: + elif isinstance(quantizer, NVFP4Quantizer) and columnwise_data is not None: k_dim = columnwise_data.shape[0] else: k_dim = tensor.shape[-1] - if tensor_offsets is None: - tensor_offsets = torch.cat( - [ - torch.zeros(1, dtype=torch.int64, device=tensor.device), - torch.cumsum(split_sizes * k_dim, dim=0), - ], - ) - return GroupedTensor( shape=(m_dim, k_dim), dtype=tensor.dtype, @@ -164,7 +171,7 @@ def _group_quantize_for_grouped_mlp( ) -> GroupedTensor: """Quantize into grouped storage.""" - if num_groups != 1 or not isinstance(quantizer, NVFP4Quantizer): + if num_groups != 1 or not isinstance(quantizer, (MXFP8Quantizer, NVFP4Quantizer)): return tex.group_quantize( tensor, quantizer, @@ -174,7 +181,7 @@ def _group_quantize_for_grouped_mlp( ) quantized = tex.quantize(tensor, quantizer) - return _wrap_single_nvfp4_as_grouped( + return _wrap_single_quantized_as_grouped( tensor, quantized, quantizer, @@ -217,7 +224,7 @@ def _group_quantize_with_amax_for_grouped_mlp( quantized = tex.nvfp4_quantize_with_amax( tensor, quantizer, rowwise_amax.view(-1)[:1], columnwise_amax.view(-1)[:1] ) - return _wrap_single_nvfp4_as_grouped( + return _wrap_single_quantized_as_grouped( tensor, quantized, quantizer, @@ -247,22 +254,29 @@ def _nvfp4_amax( return torch.cat([amax.view(-1) for amax in amaxes], dim=0) -def _nvfp4_single_tensor_from_grouped( +def _single_quantized_tensor_from_grouped( grouped: GroupedTensor, - quantizer: Optional[NVFP4Quantizer] = None, + quantizer: Optional[MXFP8Quantizer | NVFP4Quantizer] = None, *, fp4_dtype: Optional[torch.dtype] = None, -) -> NVFP4Tensor: - """Build a single NVFP4Tensor view over a one-member grouped storage.""" +) -> MXFP8Tensor | NVFP4Tensor: + """Build a single quantized tensor view over a one-member grouped storage.""" if quantizer is None: quantizer = grouped.quantizer - if not isinstance(quantizer, NVFP4Quantizer): - raise TypeError("Expected an NVFP4 GroupedTensor.") + if not isinstance(quantizer, (MXFP8Quantizer, NVFP4Quantizer)): + raise TypeError("Expected an MXFP8 or NVFP4 GroupedTensor.") shape = tuple(grouped.logical_shape) + if isinstance(quantizer, NVFP4Quantizer): + data_shape = quantizer.convert_shape_for_fp4(shape) + columnwise_shape = quantizer.convert_shape_for_fp4(quantizer.get_columnwise_shape(shape)) + else: + data_shape = shape + columnwise_shape = quantizer.get_columnwise_shape(shape) + rowwise_data = None if grouped.rowwise_data is not None: - rowwise_data = grouped.rowwise_data.view(quantizer.convert_shape_for_fp4(shape)) + rowwise_data = grouped.rowwise_data.view(data_shape) rowwise_scale_inv = None if grouped.scale_inv is not None: @@ -270,10 +284,7 @@ def _nvfp4_single_tensor_from_grouped( columnwise_data = None if grouped.columnwise_data is not None: - columnwise_shape = quantizer.get_columnwise_shape(shape) - columnwise_data = grouped.columnwise_data.view( - quantizer.convert_shape_for_fp4(columnwise_shape) - ) + columnwise_data = grouped.columnwise_data.view(columnwise_shape) columnwise_scale_inv = None if grouped.columnwise_scale_inv is not None: @@ -281,6 +292,20 @@ def _nvfp4_single_tensor_from_grouped( quantizer.get_scale_shape(shape, True) ) + if isinstance(quantizer, MXFP8Quantizer): + return MXFP8Tensor( + shape=shape, + dtype=grouped.get_dtype(), + rowwise_data=rowwise_data, + rowwise_scale_inv=rowwise_scale_inv, + columnwise_data=columnwise_data, + columnwise_scale_inv=columnwise_scale_inv, + fp8_dtype=quantizer.dtype, + quantizer=quantizer, + requires_grad=False, + with_gemm_swizzled_scales=grouped._with_gemm_swizzled_scales, + ) + return NVFP4Tensor( shape=shape, dtype=grouped.get_dtype(), @@ -335,7 +360,7 @@ def _use_tmem_post_rht_amax() -> bool: return os.environ.get("NVTE_CUTEDSL_FUSED_GROUPED_MLP_FC1_GLU_RHT_AMAX_TMEM", "0") == "1" -def _nvfp4_single_group_wgrad_gemm( +def _single_group_wgrad_gemm( grouped_x: GroupedTensor, grouped_dy: GroupedTensor, wgrad_output, @@ -343,9 +368,9 @@ def _nvfp4_single_group_wgrad_gemm( weight_shape: tuple[int, int], accumulate: bool, ) -> None: - """Run one-group NVFP4 wgrad with regular GEMM instead of grouped GEMM.""" - x_single = _nvfp4_single_tensor_from_grouped(grouped_x) - dy_single = _nvfp4_single_tensor_from_grouped(grouped_dy) + """Run one-group MXFP8/NVFP4 wgrad with regular GEMM instead of grouped GEMM.""" + x_single = _single_quantized_tensor_from_grouped(grouped_x) + dy_single = _single_quantized_tensor_from_grouped(grouped_dy) if isinstance(wgrad_output, GroupedTensor): out = wgrad_output.rowwise_data.view(1, *weight_shape)[0] else: @@ -362,6 +387,70 @@ def _nvfp4_single_group_wgrad_gemm( ) +def _single_group_fc2_gemm( + grouped_x: GroupedTensor, + grouped_weight, + input_quantizer: MXFP8Quantizer | NVFP4Quantizer, + out: torch.Tensor, + *, + single_grouped_weight: bool, + bias: Optional[torch.Tensor], + bias_scale: Optional[torch.Tensor], + dtype: torch.dtype, +) -> torch.Tensor: + """Run one-group MXFP8/NVFP4 FC2 with regular GEMM.""" + if single_grouped_weight: + weight = _single_quantized_tensor_from_grouped(grouped_weight) + else: + weight = grouped_weight[0] + + fp4_dtype = weight._fp4_dtype if isinstance(weight, NVFP4Tensor) else None + x = _single_quantized_tensor_from_grouped( + grouped_x, + input_quantizer, + fp4_dtype=fp4_dtype, + ) + general_gemm( + weight, + x, + out_dtype=dtype, + out=out, + layout="TN", + use_split_accumulator=False, + ) + + if bias is not None: + token_bias = bias.transpose(0, 1).contiguous().expand(out.shape[0], -1) + if bias_scale is not None: + out.add_(token_bias * bias_scale.view(-1, 1)) + else: + out.add_(token_bias) + return out + + +def _single_group_dgrad_gemm( + grouped_dy: GroupedTensor, + grouped_weight, + out: torch.Tensor, + *, + single_grouped_weight: bool, + dtype: torch.dtype, +) -> None: + """Run one-group MXFP8/NVFP4 dgrad with regular GEMM.""" + if single_grouped_weight: + weight = _single_quantized_tensor_from_grouped(grouped_weight) + else: + weight = grouped_weight[0] + dy = _single_quantized_tensor_from_grouped(grouped_dy) + general_gemm( + weight, + dy, + out_dtype=dtype, + out=out, + layout="NN", + ) + + def _cudnn_compute_wgrad( grouped_x: GroupedTensor, grouped_dy: GroupedTensor, @@ -541,6 +630,7 @@ def _compute_grad_params( wgrad_output = None op_label = f"Grouped MLP fused backward ({label})" if label else "Grouped MLP fused backward" weights = fc_op._get_weight_tensors() + is_dist_weight = is_distributed_weight(weights[0]) if fc_op.single_grouped_weight: w_list = [None] if ctx.weight_requires_grad: @@ -573,7 +663,9 @@ def _compute_grad_params( else: w_list = [None] * num_groups if ctx.weight_requires_grad: - if fc_op._accumulate_into_main_grad: + # Distributed weight: the GEMM produces full-sized wgrads but main_grad is sharded, + # so use a full-sized scratch buffer (the reduce-scatter below lands it in main_grad). + if fc_op._accumulate_into_main_grad and not is_dist_weight: w_list = [get_main_grad_from_param(w, op_label=op_label) for w in weights] accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) else: @@ -589,7 +681,23 @@ def _compute_grad_params( if ctx.weight_requires_grad: # Launch or defer the GEMM delay_wgrad = fc_op.wgrad_store is not None and fc_op.wgrad_store.delay_wgrad_compute() - if cudnn_wgrad_kernel_fn is not None: + if is_dist_weight and delay_wgrad: + raise RuntimeError( + "distributed-weight fused grouped-MLP requires delay_wgrad_compute=False." + ) + if ( + num_groups == 1 + and isinstance(grouped_x, (GroupedTensor, GroupedTensorStorage)) + and isinstance(grouped_dy, (GroupedTensor, GroupedTensorStorage)) + and isinstance(grouped_x.quantizer, (MXFP8Quantizer, NVFP4Quantizer)) + and isinstance(grouped_dy.quantizer, grouped_x.quantizer.__class__) + ): + gemm_fn = functools.partial( + _single_group_wgrad_gemm, + weight_shape=weight_shape, + accumulate=accumulate_into_main_grad, + ) + elif cudnn_wgrad_kernel_fn is not None: offsets = offsets if offsets.dtype == torch.int32 else offsets.to(dtype=torch.int32) gemm_fn = functools.partial( _cudnn_compute_wgrad, @@ -603,18 +711,6 @@ def _compute_grad_params( scale_view_dtype=scale_view_dtype, sf_vec_size=sf_vec_size, ) - elif ( - num_groups == 1 - and isinstance(grouped_x, GroupedTensor) - and isinstance(grouped_dy, GroupedTensor) - and isinstance(grouped_x.quantizer, NVFP4Quantizer) - and isinstance(grouped_dy.quantizer, NVFP4Quantizer) - ): - gemm_fn = functools.partial( - _nvfp4_single_group_wgrad_gemm, - weight_shape=weight_shape, - accumulate=accumulate_into_main_grad, - ) else: gemm_fn = functools.partial( general_grouped_gemm_for_grouped_tensor, @@ -627,9 +723,14 @@ def _compute_grad_params( fc_op.wgrad_store.put([grouped_x, grouped_dy, wgrad_output], gemm_fn) else: gemm_fn(grouped_x, grouped_dy, wgrad_output) + # Distributed weight: reduce-scatter the wgrads into main_grad. + # Return discarded (see finalize_weight_grads); dummy wgrads returned below. + if is_dist_weight: + finalize_weight_grads(weights, w_list) # Need to return dummy wgrads for Megatron-LM wgrad fusion if grad is already added - if fc_op._accumulate_into_main_grad: + # (wgrad fusion, or the distributed-weight reduce-scatter above) so it doesn't double-add. + if fc_op._accumulate_into_main_grad or is_dist_weight: w_list = get_dummy_wgrads_for_params(weights) elif delay_wgrad: w_list = [None] if fc_op.single_grouped_weight else [None] * num_groups @@ -878,7 +979,7 @@ def fuser_forward( prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: # Get basic operations fc1_op, activation_op, fc2_op = self.basic_ops fc1_ctx, _activation_ctx, fc2_ctx = basic_op_ctxs @@ -901,7 +1002,16 @@ def fuser_forward( # Tensor properties fc1_weight_shape = (fc1_op.out_features, fc1_op.in_features) fc2_weight_shape = (fc2_op.out_features, fc2_op.in_features) - input_ = input_.reshape(-1, fc1_weight_shape[1]) + if isinstance(input_, GroupedTensor): + # GroupedTensor forbids reshape and is already in the canonical + # (total_tokens, in_features) layout; just validate the shape. + if input_.dim() != 2 or input_.size(-1) != fc1_weight_shape[1]: + raise ValueError( + "GroupedTensor input must have shape (total_tokens, " + f"{fc1_weight_shape[1]}), but got {tuple(input_.size())}." + ) + else: + input_ = input_.reshape(-1, fc1_weight_shape[1]) in_shape = list(input_.size()) if in_shape[0] % 128 != 0: raise ValueError(f"Unsupported input shape for fused grouped MLP ({in_shape=}).") @@ -909,6 +1019,19 @@ def fuser_forward( num_groups = fc1_op.num_groups fc1_weight_param = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0 fc2_weight_param = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0 + + # Distributed weight: expert weights are sharded 1/N along out_features; the fused kernels + # read the full shape, so all-gather the full weight first. A plain weight is a no-op. + fc1_is_dist = is_distributed_weight(fc1_weight_param) + fc2_is_dist = is_distributed_weight(fc2_weight_param) + assert fc1_is_dist == fc2_is_dist, "FC1/FC2 must share one distributed-weight group." + if fc1_is_dist: + assert ( + not fc1_op.single_grouped_weight and not fc2_op.single_grouped_weight + ), "distributed-weight fused grouped-MLP requires single_grouped_weight=False." + assert fc1_op.weight0.is_routed_expert and fc1_op.weight0.weight_list is not None + assert fc2_op.weight0.is_routed_expert and fc2_op.weight0.weight_list is not None + device = fc1_weight_param.device if torch.is_autocast_enabled(): dtype = torch.get_autocast_dtype("cuda") @@ -944,24 +1067,79 @@ def fuser_forward( if int(split_sizes.numel()) != num_groups: raise ValueError(f"Expected {num_groups} splits, but got {int(split_sizes.numel())}.") - # Prepare split metadata - split_sizes, ( - split_points, - base_split_offsets, - fc1_x_tensor_offsets, - fc2_x_tensor_offsets, - fc2_out_tensor_offsets, - ) = tex.splits_to_offsets_multi( - split_sizes, - device, - strides=[1, 1, fc1_weight_shape[1], fc2_weight_shape[1], fc2_weight_shape[0]], - include_leading_zero=[False, True, True, True, True], - dtypes=[torch.int32, torch.int64, torch.int64, torch.int64, torch.int64], - bulk_allocate=True, - ) - # Extract per-row activation probabilities from the middle op. scales = basic_op_extra_inputs[1][0] + unit_activation_scale = bool( + getattr(self.basic_ops[1], "_grouped_mlp_unit_activation_scale", False) + ) + if unit_activation_scale and num_groups != 1: + unit_activation_scale = False + + activation_kernel = self.grouped_gemm_activation_kernel() + supports_single_group_runtime_offsets = ( + _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)) + ) + + # Shared experts have one dense group and all optimized kernels derive M + # from their runtime tensor shapes. Reuse the caller-owned split tensor + # as the ignored cuDNN padded-offset argument and omit tensor offsets + # entirely. This avoids both splits_to_offsets and cached CUDA pointers. + # Older cuDNN frontends do not expose this specialization, so use the + # live generic offset calculation rather than caching CUDA metadata. + use_offsetless_metadata = ( + num_groups == 1 + and unit_activation_scale + and isinstance(fc1_input_quantizer, MXFP8Quantizer) + and supports_single_group_runtime_offsets + ) + + # Prepare split metadata + if use_offsetless_metadata: + # cuDNN requires an int32 tensor descriptor, although its + # use_single_group_runtime_offsets specialization never loads the + # pointer. This view aliases the live caller-owned [M] int64 tensor, + # has value M, and requires no allocation or CUDA kernel. + split_points = split_sizes.view(torch.int32)[:1] + # Backward saves this slot for the generic path. The optimized + # shared-expert path never consumes it. + base_split_offsets = split_sizes + fc1_x_tensor_offsets = None + fc1_out_tensor_offsets = None + fc2_x_tensor_offsets = None + fc2_out_tensor_offsets = None + else: + # Bulk-allocate every grouped-tensor offset the forward and backward + # passes need, so the backward can reuse them from the context + # instead of recomputing offsets per GEMM. + split_sizes, ( + split_points, + base_split_offsets, + fc1_x_tensor_offsets, + fc1_out_tensor_offsets, + fc2_x_tensor_offsets, + fc2_out_tensor_offsets, + ) = tex.splits_to_offsets_multi( + split_sizes, + device, + strides=[ + 1, + 1, + fc1_weight_shape[1], + fc1_weight_shape[0], + fc2_weight_shape[1], + fc2_weight_shape[0], + ], + include_leading_zero=[False, True, True, True, True, True], + dtypes=[ + torch.int32, + torch.int64, + torch.int64, + torch.int64, + torch.int64, + torch.int64, + ], + bulk_allocate=True, + ) # Prepare FC1 grouped weight tensor for fused kernels. # - single_grouped_weight=True: op.weight is already a GroupedTensor @@ -973,8 +1151,13 @@ def fuser_forward( "FC1 expected GroupedTensor weight with single_grouped_weight=True." ) if fc1_op.weight.quantizer is not None: - fc1_weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) - fc1_op.weight.quantizer = fc1_weight_quantizer + # Update param quantizer + # Note: Only add usages, never drop. The same weight may be used in multiple steps. + fc1_weight_quantizer = fc1_op.weight.quantizer + fc1_weight_quantizer.set_usage( + rowwise=True, + columnwise=True if input_requires_grad else None, + ) grouped_fc1_weight = fc1_op.weight else: if fc1_op.weight.rowwise_data is None: @@ -997,6 +1180,8 @@ def fuser_forward( else: quantized_fc1_weights.append(weight) grouped_fc1_weight = quantized_fc1_weights + if fc1_is_dist: + grouped_fc1_weight = materialize_weight_for_forward(grouped_fc1_weight) # Prepare FC2 grouped weight tensor for fused kernels. if fc2_op.single_grouped_weight: @@ -1005,8 +1190,11 @@ def fuser_forward( "FC2 expected GroupedTensor weight with single_grouped_weight=True." ) if fc2_op.weight.quantizer is not None: - fc2_weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) - fc2_op.weight.quantizer = fc2_weight_quantizer + fc2_weight_quantizer = fc2_op.weight.quantizer + fc2_weight_quantizer.set_usage( + rowwise=True, + columnwise=True if requires_grad else None, + ) grouped_fc2_weight = fc2_op.weight else: if fc2_op.weight.rowwise_data is None: @@ -1035,51 +1223,26 @@ def fuser_forward( grouped_fc1_weight, "_with_gemm_swizzled_scales" ): grouped_fc1_weight._with_gemm_swizzled_scales = False - if isinstance(grouped_fc2_weight, GroupedTensor) and not hasattr( - grouped_fc2_weight, "_with_gemm_swizzled_scales" - ): - grouped_fc2_weight._with_gemm_swizzled_scales = False # Group-quantize input tensor and convert dtypes if needed - fc1_input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) + fc1_input_quantizer.set_usage( + rowwise=True, + columnwise=weight_requires_grad, + ) fc1_input_quantizer.optimize_for_gemm = True fc1_input_quantizer.internal = True - input_quantizer = getattr(input_, "quantizer", None) - if isinstance(input_, GroupedTensor) and ( - isinstance(fc1_input_quantizer, MXFP8Quantizer) - and isinstance(input_quantizer, MXFP8Quantizer) - or isinstance(fc1_input_quantizer, NVFP4Quantizer) - and isinstance(input_quantizer, NVFP4Quantizer) - ): - # GroupedTensor is a torch.Tensor subclass, so the CPU offload - # infrastructure's prepare_for_saving treats it as a plain tensor - # and does not decompose it into its component data tensors. By - # repacking into a GroupedTensorStorage (not a torch.Tensor), we - # ensure the fuser's prepare_for_saving call correctly decomposes - # the activation before save_for_backward. - grouped_fc1_x = GroupedTensorStorage( - shape=input_.logical_shape, - dtype=input_.fake_dtype, - num_tensors=input_.num_tensors, - shapes=input_.tensor_shapes, - quantizer=input_.quantizer, - data=input_.rowwise_data, - columnwise_data=input_.columnwise_data, - scale_inv=input_.scale_inv, - columnwise_scale_inv=input_.columnwise_scale_inv, - amax=input_.amax, - columnwise_amax=input_.columnwise_amax, - scale=input_.scale, - first_dims=input_.first_dims, - last_dims=input_.last_dims, - tensor_offsets=input_.tensor_offsets, - offsets=input_.offsets, - scale_inv_offsets=input_.scale_inv_offsets, - columnwise_scale_inv_offsets=input_.columnwise_scale_inv_offsets, - with_gemm_swizzled_scales=input_._with_gemm_swizzled_scales, - row_scaled_nvfp4=input_.row_scaled_nvfp4, - nvfp4_use_4over6=input_.nvfp4_use_4over6, - nvfp4_e4m3_max=input_.nvfp4_e4m3_max, + if isinstance(input_, GroupedTensor): + # Input arrived already quantized (e.g. FP8 token dispatch): reuse its rowwise data + # for the GEMM and let the helper supply whatever else the GEMMs need. An input that + # is already GEMM-ready in both directions passes through untouched. + grouped_fc1_x = input_.copy() + tex.group_requantize_inplace( + grouped_fc1_x, + fc1_input_quantizer, + num_groups, + split_sizes, + TE_DType[dtype], + tensor_offsets=fc1_x_tensor_offsets, ) else: fc1_x = maybe_dequantize(input_, dtype) @@ -1155,9 +1318,11 @@ def fuser_forward( fc2_bias_packed = _pack_grouped_linear_bias_for_cudnn(fc2_op) fc1_d_dtype = torch.bfloat16 if use_nvfp4 else torch.float8_e4m3fn - fc1_prob_tensor = ( - scales.detach().to(dtype=torch.float32 if use_nvfp4 else dtype).reshape(-1, 1, 1) - ) + fc1_prob_tensor = None + if not unit_activation_scale: + fc1_prob_tensor = ( + scales.detach().to(dtype=torch.float32 if use_nvfp4 else dtype).reshape(-1, 1, 1) + ) fc1_norm_const_tensor = None if use_nvfp4 else norm_const_tensor if use_nvfp4: nvfp4_fp4_max = 6.0 @@ -1217,6 +1382,8 @@ def fuser_forward( else: fc1_activation_kwargs["norm_const_tensor"] = fc1_norm_const_tensor fc1_activation_kwargs["discrete_col_sfd"] = not use_nvfp4 + if supports_single_group_runtime_offsets: + fc1_activation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 if self._pass_geglu_runtime_params: fc1_activation_kwargs.update( linear_offset=self._cudnn_linear_offset, @@ -1228,16 +1395,30 @@ def fuser_forward( if fc1_op.single_grouped_weight: # Clone and swizzle scales for GEMM. fc1_weight_for_gemm = grouped_fc1_weight.copy() - tex.grouped_swizzle_for_gemm(fc1_weight_for_gemm, rowwise=True, columnwise=False) + use_single_group_weight_swizzle = num_groups == 1 + if use_single_group_weight_swizzle: + fc1_weight_single = _single_quantized_tensor_from_grouped(fc1_weight_for_gemm) + fc1_weight_single._columnwise_data = None + fc1_weight_single._columnwise_scale_inv = None + tex.swizzle_scales_for_gemm_(fc1_weight_single) + fc1_w_data = fc1_weight_single._rowwise_data + fc1_w_scales = fc1_weight_single._rowwise_scale_inv + else: + tex.grouped_swizzle_for_gemm( + fc1_weight_for_gemm, + rowwise=True, + columnwise=False, + ) + fc1_w_data = fc1_weight_for_gemm.rowwise_data + fc1_w_scales = fc1_weight_for_gemm.scale_inv # Pack weight tensors for stacked kernel # Data actual shape: (num_groups, n, k) # Data logical shape: (n, k, num_groups) - fc1_w_data = fc1_weight_for_gemm.rowwise_data fc1_w_data = fc1_w_data.view(dtype=data_dtype) fc1_w_data = fc1_w_data.view(num_groups, fc1_weight_shape[0], fc1_weight_k) fc1_w_data = fc1_w_data.permute(1, 2, 0) - fc1_w_scales = fc1_weight_for_gemm.scale_inv.view(dtype=scale_view_dtype) + fc1_w_scales = fc1_w_scales.view(dtype=scale_view_dtype) fc1_w_scales = fc1_w_scales.view( num_groups, ceil_div(fc1_weight_shape[0], 128), @@ -1251,25 +1432,68 @@ def fuser_forward( fc1_activation_kwargs["b_tensor"] = fc1_w_data fc1_activation_kwargs["sfb_tensor"] = fc1_w_scales else: - # Discrete-weight kernel: per-expert data/scale pointers - fc1_b_ptrs, fc1_sfb_ptrs, _fc1_sfb_buffer = ( - tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( - [w._rowwise_data for w in grouped_fc1_weight], - [w._rowwise_scale_inv for w in grouped_fc1_weight], - "nvfp4" if use_nvfp4 else "mxfp8_rowwise", - device, + use_single_discrete_weight = num_groups == 1 + if use_single_discrete_weight: + fc1_weight_single = grouped_fc1_weight[0] + original_rowwise_scale = fc1_weight_single._rowwise_scale_inv + original_columnwise_data = fc1_weight_single._columnwise_data + original_columnwise_scale = fc1_weight_single._columnwise_scale_inv + original_swizzled = fc1_weight_single._with_gemm_swizzled_scales + fc1_weight_single._columnwise_data = None + fc1_weight_single._columnwise_scale_inv = None + tex.swizzle_scales_for_gemm_(fc1_weight_single) + swizzled_rowwise_scale = fc1_weight_single._rowwise_scale_inv + fc1_weight_single._rowwise_scale_inv = original_rowwise_scale + fc1_weight_single._columnwise_data = original_columnwise_data + fc1_weight_single._columnwise_scale_inv = original_columnwise_scale + fc1_weight_single._with_gemm_swizzled_scales = original_swizzled + + fc1_w_data = fc1_weight_single._rowwise_data.view(dtype=data_dtype) + fc1_w_data = fc1_w_data.view( + 1, + fc1_weight_shape[0], + fc1_weight_k, ) - ) - fc1_activation_kwargs["b_ptrs"] = fc1_b_ptrs - fc1_activation_kwargs["sfb_ptrs"] = fc1_sfb_ptrs - fc1_activation_kwargs["n"] = fc1_weight_shape[0] - fc1_activation_kwargs["b_dtype"] = data_dtype - fc1_activation_kwargs["b_major"] = "k" + fc1_w_data = fc1_w_data.permute(1, 2, 0) + fc1_w_scales = swizzled_rowwise_scale.view(dtype=scale_view_dtype) + fc1_w_scales = fc1_w_scales.view( + 1, + ceil_div(fc1_weight_shape[0], 128), + ceil_div(fc1_weight_shape[1], k_sf_divisor), + 32, + 4, + 4, + ) + fc1_w_scales = fc1_w_scales.permute(3, 4, 1, 5, 2, 0) + fc1_activation_kwargs["b_tensor"] = fc1_w_data + fc1_activation_kwargs["sfb_tensor"] = fc1_w_scales + else: + # Discrete-weight kernel: per-expert data/scale pointers + fc1_b_ptrs, fc1_sfb_ptrs, _fc1_sfb_buffer = ( + tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( + [w._rowwise_data for w in grouped_fc1_weight], + [w._rowwise_scale_inv for w in grouped_fc1_weight], + "nvfp4" if use_nvfp4 else "mxfp8_rowwise", + device, + ) + ) + fc1_activation_kwargs["b_ptrs"] = fc1_b_ptrs + fc1_activation_kwargs["sfb_ptrs"] = fc1_sfb_ptrs + fc1_activation_kwargs["n"] = fc1_weight_shape[0] + fc1_activation_kwargs["b_dtype"] = data_dtype + fc1_activation_kwargs["b_major"] = "k" if use_fc1_act_hadamard: fc1_kernel_out = self.grouped_gemm_act_hadamard_kernel()(**fc1_activation_kwargs) else: - fc1_kernel_out = self.grouped_gemm_activation_kernel()(**fc1_activation_kwargs) + fc1_kernel_out = activation_kernel(**fc1_activation_kwargs) + + if fc2_is_dist: + grouped_fc2_weight = materialize_weight_for_forward(grouped_fc2_weight) + if isinstance(grouped_fc2_weight, GroupedTensor) and not hasattr( + grouped_fc2_weight, "_with_gemm_swizzled_scales" + ): + grouped_fc2_weight._with_gemm_swizzled_scales = False # Unpack kernel outputs # Note: Fused kernel outputs tensors with non-contiguous @@ -1326,31 +1550,16 @@ def fuser_forward( and grouped_fc2_x.columnwise_data is not None and grouped_fc2_x.columnwise_scale_inv is not None ): - if fc2_op.single_grouped_weight: - fc2_w_single = grouped_fc2_weight.split_into_quantized_tensors()[0] - else: - fc2_w_single = grouped_fc2_weight[0] - fc2_x_single = _nvfp4_single_tensor_from_grouped( + fc2_out_buf = _single_group_fc2_gemm( grouped_fc2_x, + grouped_fc2_weight, fc2_input_quantizer, - fp4_dtype=fc2_w_single._fp4_dtype, - ) - general_gemm( - fc2_w_single, - fc2_x_single, - out_dtype=dtype, - out=fc2_out_buf, - layout="TN", - use_split_accumulator=False, + fc2_out_buf, + single_grouped_weight=fc2_op.single_grouped_weight, + bias=fc2_bias_packed, + bias_scale=fc2_scales, + dtype=dtype, ) - if fc2_bias_packed is not None: - token_bias = ( - fc2_bias_packed.transpose(0, 1).contiguous().expand(in_shape[0], -1) - ) - if fc2_scales is not None: - fc2_out_buf += token_bias * fc2_scales.view(-1, 1) - else: - fc2_out_buf += token_bias else: fc2_out_grouped = GroupedTensorStorage( shape=(in_shape[0], fc2_weight_shape[0]), @@ -1395,77 +1604,106 @@ def fuser_forward( with_gemm_swizzled_scales=True, ) - fc2_scales_tensor = ( - fc2_scales.detach().to(dtype=torch.float32).reshape(-1, 1, 1) - if fc2_scales is not None - else torch.ones((in_shape[0], 1, 1), dtype=torch.float32, device=device) - ) - fc2_quant_kwargs = { - "a_tensor": fc1_kernel_out["d_tensor"], - "sfa_tensor": fc1_kernel_out["sfd_row_tensor"], - "padded_offsets": split_points, - "alpha_tensor": alpha_tensor, - "bias_tensor": fc2_bias_packed, - "norm_const_tensor": None, - "prob_tensor": fc2_scales_tensor, - "acc_dtype": torch.float32, - "d_dtype": dtype, - "cd_major": "n", - "sf_vec_size": MXFP8_BLOCK_SCALING_SIZE, - "current_stream": current_stream, - "use_dynamic_sched": True, - } - - if fc2_op.single_grouped_weight: - # Clone and swizzle scales for GEMM (original stays unmodified for save_for_backward) - fc2_weight_for_gemm = grouped_fc2_weight.copy() - tex.grouped_swizzle_for_gemm(fc2_weight_for_gemm, rowwise=True, columnwise=False) - - fc2_w_data = fc2_weight_for_gemm.rowwise_data - fc2_w_data = fc2_w_data.view(dtype=torch.float8_e4m3fn) - fc2_w_data = fc2_w_data.view(num_groups, fc2_weight_shape[0], fc2_weight_shape[1]) - fc2_w_data = fc2_w_data.permute(1, 2, 0) - - fc2_w_scales = fc2_weight_for_gemm.scale_inv.view(dtype=torch.float8_e8m0fnu) - fc2_w_scales = fc2_w_scales.view( - num_groups, - ceil_div(fc2_weight_shape[0], 128), - ceil_div(fc2_weight_shape[1], 128), - MXFP8_BLOCK_SCALING_SIZE, - 4, - 4, + use_single_group_dense_fc2 = num_groups == 1 + fc2_out_buf = validate_or_alloc_output(output_buffer, fc2_out_shape, dtype, device) + if use_single_group_dense_fc2: + fc2_out = _single_group_fc2_gemm( + grouped_fc2_x, + grouped_fc2_weight, + fc2_input_quantizer, + fc2_out_buf, + single_grouped_weight=fc2_op.single_grouped_weight, + bias=fc2_bias_packed, + bias_scale=fc2_scales, + dtype=dtype, ) - fc2_w_scales = fc2_w_scales.permute(3, 4, 1, 5, 2, 0) - fc2_quant_kwargs["b_tensor"] = fc2_w_data - fc2_quant_kwargs["sfb_tensor"] = fc2_w_scales else: - fc2_b_ptrs, fc2_sfb_ptrs, _fc2_sfb_buffer = ( - tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( - [w._rowwise_data for w in grouped_fc2_weight], - [w._rowwise_scale_inv for w in grouped_fc2_weight], - "nvfp4" if use_nvfp4 else "mxfp8_rowwise", - device, + cudnn_supports_optional_prob = _cudnn_frontend_version_at_least("1.27.0") + fc2_scales_tensor = ( + None + if cudnn_supports_optional_prob + else torch.ones((in_shape[0], 1, 1), dtype=torch.float32, device=device) + ) + if fc2_scales is not None: + fc2_scales_tensor = ( + fc2_scales.detach().to(dtype=torch.float32).reshape(-1, 1, 1) + ) + fc2_quant_kwargs = { + "a_tensor": fc1_kernel_out["d_tensor"], + "sfa_tensor": fc1_kernel_out["sfd_row_tensor"], + "padded_offsets": split_points, + "alpha_tensor": alpha_tensor, + "bias_tensor": fc2_bias_packed, + "norm_const_tensor": None, + "prob_tensor": fc2_scales_tensor, + "acc_dtype": torch.float32, + "d_dtype": dtype, + "cd_major": "n", + "sf_vec_size": MXFP8_BLOCK_SCALING_SIZE, + "current_stream": current_stream, + "use_dynamic_sched": True, + } + fc2_quant_kernel = self.grouped_gemm_quant_kernel() + if supports_single_group_runtime_offsets: + fc2_quant_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 + + if fc2_op.single_grouped_weight: + # Clone and swizzle scales for GEMM + # (original stays unmodified for save_for_backward). + fc2_weight_for_gemm = grouped_fc2_weight.copy() + tex.grouped_swizzle_for_gemm( + fc2_weight_for_gemm, + rowwise=True, + columnwise=False, + ) + + fc2_w_data = fc2_weight_for_gemm.rowwise_data + fc2_w_data = fc2_w_data.view(dtype=torch.float8_e4m3fn) + fc2_w_data = fc2_w_data.view( + num_groups, + fc2_weight_shape[0], + fc2_weight_shape[1], + ) + fc2_w_data = fc2_w_data.permute(1, 2, 0) + + fc2_w_scales = fc2_weight_for_gemm.scale_inv.view(dtype=torch.float8_e8m0fnu) + fc2_w_scales = fc2_w_scales.view( + num_groups, + ceil_div(fc2_weight_shape[0], 128), + ceil_div(fc2_weight_shape[1], 128), + MXFP8_BLOCK_SCALING_SIZE, + 4, + 4, + ) + fc2_w_scales = fc2_w_scales.permute(3, 4, 1, 5, 2, 0) + fc2_quant_kwargs["b_tensor"] = fc2_w_data + fc2_quant_kwargs["sfb_tensor"] = fc2_w_scales + else: + fc2_b_ptrs, fc2_sfb_ptrs, _fc2_sfb_buffer = ( + tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( + [w._rowwise_data for w in grouped_fc2_weight], + [w._rowwise_scale_inv for w in grouped_fc2_weight], + "mxfp8_rowwise", + device, + ) ) + fc2_quant_kwargs["b_ptrs"] = fc2_b_ptrs + fc2_quant_kwargs["sfb_ptrs"] = fc2_sfb_ptrs + fc2_quant_kwargs["n"] = fc2_weight_shape[0] + fc2_quant_kwargs["b_dtype"] = torch.float8_e4m3fn + fc2_quant_kwargs["b_major"] = "k" + + fc2_quant_kwargs["d_tensor"] = fc2_out_buf.as_strided( + (in_shape[0], fc2_weight_shape[0], 1), + (fc2_weight_shape[0], 1, in_shape[0] * fc2_weight_shape[0]), ) - fc2_quant_kwargs["b_ptrs"] = fc2_b_ptrs - fc2_quant_kwargs["sfb_ptrs"] = fc2_sfb_ptrs - fc2_quant_kwargs["n"] = fc2_weight_shape[0] - fc2_quant_kwargs["b_dtype"] = torch.float8_e4m3fn - fc2_quant_kwargs["b_major"] = "k" - - # Always allocate the output (the caller's buffer if provided, else a fresh one) and - # pass it as the kernel's d_tensor, so the kernel writes in place and the call is uniform. - output_buffer = validate_or_alloc_output(output_buffer, fc2_out_shape, dtype, device) - fc2_quant_kwargs["d_tensor"] = output_buffer.as_strided( - (in_shape[0], fc2_weight_shape[0], 1), - (fc2_weight_shape[0], 1, in_shape[0] * fc2_weight_shape[0]), - ) - self.grouped_gemm_quant_kernel()(**fc2_quant_kwargs) - fc2_out = output_buffer + fc2_quant_kernel(**fc2_quant_kwargs) + fc2_out = fc2_out_buf # Save state for backward pass if requires_grad: - mark_grouped_tensor(grouped_fc1_x, activation_in, scales, grouped_fc2_x) + saved_fc1_x = grouped_fc1_x + mark_grouped_tensor(saved_fc1_x, activation_in, scales, grouped_fc2_x) activation_op = self.basic_ops[1] cpu_offloading = is_cpu_offload_enabled() activation_is_srelu = isinstance(activation_op, ScaledSReLU) @@ -1484,14 +1722,14 @@ def fuser_forward( # MXFP8 wgrad only needs columnwise tiles. NVFP4 generic GEMM fallbacks # need the full grouped tensor state, including rowwise data and amax. if not use_nvfp4: - for grouped_fc_x in (grouped_fc1_x, saved_grouped_fc2_x): - if grouped_fc_x is not None: + for grouped_fc_x in (saved_fc1_x, saved_grouped_fc2_x): + if isinstance(grouped_fc_x, (GroupedTensor, GroupedTensorStorage)): grouped_fc_x.rowwise_data = None grouped_fc_x.scale_inv = None if cpu_offloading: activation_tensors = [ - t for t in (grouped_fc1_x, activation_in, saved_grouped_fc2_x) if t is not None + t for t in (saved_fc1_x, activation_in, saved_grouped_fc2_x) if t is not None ] start_offload(*activation_tensors) mark_activation_offload(*activation_tensors) @@ -1504,11 +1742,20 @@ def fuser_forward( fc2_weight_tensors = ( [grouped_fc2_weight] if fc2_op.single_grouped_weight else grouped_fc2_weight ) + # Save the joint op's internal layout; distributed weights save the small shards. + if fc1_is_dist: + fc1_weight_tensors = fc1_weights + if fc2_is_dist: + fc2_weight_tensors = fc2_weights fc1_ctx.save_for_backward( split_sizes, base_split_offsets, split_points, - grouped_fc1_x, + fc1_x_tensor_offsets, + fc1_out_tensor_offsets, + fc2_x_tensor_offsets, + fc2_out_tensor_offsets, + saved_fc1_x, *fc1_weight_tensors, activation_in, scales, @@ -1521,6 +1768,7 @@ def fuser_forward( fc1_ctx.dtype = dtype fc1_ctx.input_requires_grad = input_requires_grad fc1_ctx.weight_requires_grad = weight_requires_grad + fc1_ctx.unit_activation_scale = unit_activation_scale fc2_ctx.input_quantizers = [fc2_input_quantizer] fc2_ctx.grad_output_quantizers = [fc2_grad_output_quantizer] @@ -1550,7 +1798,16 @@ def fuser_backward( # Tensor properties fc1_weight_shape = (fc1_op.out_features, fc1_op.in_features) fc2_weight_shape = (fc2_op.out_features, fc2_op.in_features) - grad_output = grad_output.reshape(-1, fc2_weight_shape[0]) + if isinstance(grad_output, GroupedTensor): + # GroupedTensor forbids reshape and is already in the canonical + # (total_tokens, out_features) layout; just validate the shape. + if grad_output.dim() != 2 or grad_output.size(-1) != fc2_weight_shape[0]: + raise ValueError( + "GroupedTensor grad output must have shape (total_tokens, " + f"{fc2_weight_shape[0]}), but got {tuple(grad_output.size())}." + ) + else: + grad_output = grad_output.reshape(-1, fc2_weight_shape[0]) out_shape = list(grad_output.size()) num_groups = fc1_op.num_groups fc1_weight_param = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0 @@ -1560,12 +1817,22 @@ def fuser_backward( # Saved tensors from the joint forward. # Layout: [split_sizes, base_split_offsets, split_points, + # fc1_x_tensor_offsets, fc1_out_tensor_offsets, + # fc2_x_tensor_offsets, fc2_out_tensor_offsets, # grouped_fc1_x, *fc1_weights, # activation_in, scales, # grouped_fc2_x, *fc2_weights] saved_tensors = fc1_ctx.saved_tensors - split_sizes, base_split_offsets, split_points = saved_tensors[:3] - saved_tensors = saved_tensors[3:] + ( + split_sizes, + base_split_offsets, + split_points, + fc1_x_tensor_offsets, + fc1_out_tensor_offsets, + fc2_x_tensor_offsets, + fc2_out_tensor_offsets, + ) = saved_tensors[:7] + saved_tensors = saved_tensors[7:] grouped_fc1_x, saved_tensors = saved_tensors[0], saved_tensors[1:] if fc1_op.single_grouped_weight: grouped_fc1_weight, saved_tensors = saved_tensors[0], saved_tensors[1:] @@ -1610,20 +1877,28 @@ def fuser_backward( output_fc2_dbias = fc2_op.has_bias fc2_dbias_packed = None fc2_dy = None - grad_output_quantizer = getattr(grad_output, "quantizer", None) - fc2_grad_output_quantizer_matches = ( - isinstance(fc2_grad_output_quantizer, MXFP8Quantizer) - and isinstance(grad_output_quantizer, MXFP8Quantizer) - ) or ( - isinstance(fc2_grad_output_quantizer, NVFP4Quantizer) - and isinstance(grad_output_quantizer, NVFP4Quantizer) - ) - if ( - not output_fc2_dbias - and isinstance(grad_output, GroupedTensor) - and fc2_grad_output_quantizer_matches - ): - grouped_fc2_dy = grad_output + if isinstance(grad_output, GroupedTensor): + # Grad output arrived already quantized (e.g. FP8 token dispatch): reuse its rowwise + # data for the dgrad GEMM. Bias grads are reduced from the dequantized grad, which is + # only materialized when one is needed. A grad that is already GEMM-ready in both + # directions passes through untouched. + grouped_fc2_dy = grad_output.copy() + fc2_dy = tex.group_requantize_inplace( + grouped_fc2_dy, + fc2_grad_output_quantizer, + num_groups, + split_sizes, + TE_DType[dtype], + tensor_offsets=fc2_out_tensor_offsets, + return_dequantized=output_fc2_dbias or scale_bias, + ) + if output_fc2_dbias and not scale_bias: + # This path has no quantize kernel to fuse dbias into, and the consumer below + # has no fallback, so reduce it here. + fc2_dbias_packed = compute_grouped_dbias(fc2_dy, base_split_offsets, num_groups) + # scale_bias is the only later consumer of the dequantized grad; drop it so the + # buffer is freed rather than held until backward ends. + fc2_dy = None else: fc2_dy = maybe_dequantize(grad_output, dtype) if output_fc2_dbias and not scale_bias: @@ -1632,6 +1907,7 @@ def fuser_backward( fc2_grad_output_quantizer, num_groups, split_sizes, + tensor_offsets=fc2_out_tensor_offsets, ) else: grouped_fc2_dy = _group_quantize_for_grouped_mlp( @@ -1639,7 +1915,7 @@ def fuser_backward( fc2_grad_output_quantizer, num_groups, split_sizes, - tensor_offsets=base_split_offsets * fc2_weight_shape[0], + tensor_offsets=fc2_out_tensor_offsets, ) use_nvfp4 = ( @@ -1705,9 +1981,14 @@ def fuser_backward( norm_const_tensor = get_cached_ones_tensor(1, torch.float32, device) current_stream = torch.cuda.current_stream().cuda_stream - scales_f32 = scales.detach().to(dtype=torch.float32) - scales_tensor = scales_f32.reshape(-1, 1, 1) - dscales_tensor = torch.zeros_like(scales_tensor) + unit_activation_scale = bool(getattr(fc1_ctx, "unit_activation_scale", False)) + scales_f32 = None + scales_tensor = None + dscales_tensor = None + if not unit_activation_scale: + scales_f32 = scales.detach().to(dtype=torch.float32) + scales_tensor = scales_f32.reshape(-1, 1, 1) + dscales_tensor = torch.zeros_like(scales_tensor) fc2_d_dtype = torch.bfloat16 if use_nvfp4 else torch.float8_e4m3fn if use_nvfp4: @@ -1754,6 +2035,9 @@ def fuser_backward( "discrete_col_sfd": not use_nvfp4, "use_dynamic_sched": True, } + dactivation_kernel = self.grouped_gemm_dactivation_kernel() + if _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)): + fc2_dactivation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 if self._cudnn_dact_func is not None: fc2_dactivation_kwargs["beta_tensor"] = fc2_beta_tensor fc2_dactivation_kwargs["act_func"] = self._cudnn_dact_func @@ -1767,6 +2051,10 @@ def fuser_backward( glu_clamp_min=self._cudnn_glu_clamp_min, ) + fc2_leader = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0 + if is_distributed_weight(fc2_leader): + grouped_fc2_weight = materialize_weight_for_backward(fc2_leader) + if fc2_op.single_grouped_weight: # Clone and swizzle scales for GEMM fc2_weight_for_gemm = grouped_fc2_weight.copy() @@ -1796,21 +2084,63 @@ def fuser_backward( fc2_dactivation_kwargs["b_tensor"] = fc2_w_data fc2_dactivation_kwargs["sfb_tensor"] = fc2_w_scales else: - fc2_b_ptrs, fc2_sfb_ptrs, _fc2_sfb_buffer = ( - tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( - [w._columnwise_data for w in grouped_fc2_weight], - [w._columnwise_scale_inv for w in grouped_fc2_weight], - "nvfp4" if use_nvfp4 else "mxfp8_columnwise", - device, + use_single_discrete_weight = num_groups == 1 + if use_single_discrete_weight: + fc2_weight_single = grouped_fc2_weight[0] + original_rowwise_data = fc2_weight_single._rowwise_data + original_rowwise_scale = fc2_weight_single._rowwise_scale_inv + original_columnwise_scale = fc2_weight_single._columnwise_scale_inv + original_swizzled = fc2_weight_single._with_gemm_swizzled_scales + fc2_weight_single._rowwise_data = None + fc2_weight_single._rowwise_scale_inv = None + tex.swizzle_scales_for_gemm_(fc2_weight_single) + swizzled_columnwise_scale = fc2_weight_single._columnwise_scale_inv + fc2_weight_single._rowwise_data = original_rowwise_data + fc2_weight_single._rowwise_scale_inv = original_rowwise_scale + fc2_weight_single._columnwise_scale_inv = original_columnwise_scale + fc2_weight_single._with_gemm_swizzled_scales = original_swizzled + + fc2_w_data = fc2_weight_single._columnwise_data.view(dtype=data_dtype) + fc2_w_data = fc2_w_data.view( + 1, + fc2_weight_shape[0], + fc2_weight_k, ) - ) - fc2_dactivation_kwargs["b_ptrs"] = fc2_b_ptrs - fc2_dactivation_kwargs["sfb_ptrs"] = fc2_sfb_ptrs - fc2_dactivation_kwargs["n"] = fc2_weight_shape[1] - fc2_dactivation_kwargs["b_dtype"] = data_dtype - fc2_dactivation_kwargs["b_major"] = "k" if use_nvfp4 else "n" + fc2_w_data = ( + fc2_w_data.permute(1, 2, 0) if use_nvfp4 else fc2_w_data.permute(2, 1, 0) + ) + fc2_w_scales = swizzled_columnwise_scale.view(dtype=scale_view_dtype) + fc2_w_scales = fc2_w_scales.view( + 1, + ceil_div(fc2_weight_shape[1], k_sf_divisor), + ceil_div(fc2_weight_shape[0], 128), + 32, + 4, + 4, + ) + fc2_w_scales = ( + fc2_w_scales.permute(3, 4, 2, 5, 1, 0) + if use_nvfp4 + else fc2_w_scales.permute(3, 4, 1, 5, 2, 0) + ) + fc2_dactivation_kwargs["b_tensor"] = fc2_w_data + fc2_dactivation_kwargs["sfb_tensor"] = fc2_w_scales + else: + fc2_b_ptrs, fc2_sfb_ptrs, _fc2_sfb_buffer = ( + tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( + [w._columnwise_data for w in grouped_fc2_weight], + [w._columnwise_scale_inv for w in grouped_fc2_weight], + "nvfp4" if use_nvfp4 else "mxfp8_columnwise", + device, + ) + ) + fc2_dactivation_kwargs["b_ptrs"] = fc2_b_ptrs + fc2_dactivation_kwargs["sfb_ptrs"] = fc2_sfb_ptrs + fc2_dactivation_kwargs["n"] = fc2_weight_shape[1] + fc2_dactivation_kwargs["b_dtype"] = data_dtype + fc2_dactivation_kwargs["b_major"] = "k" if use_nvfp4 else "n" - fc2_dgrad_kernel_out = self.grouped_gemm_dactivation_kernel()(**fc2_dactivation_kwargs) + fc2_dgrad_kernel_out = dactivation_kernel(**fc2_dactivation_kwargs) if use_nvfp4: fc1_dy_bf16 = fc2_dgrad_kernel_out["d_row_tensor"] @@ -1833,7 +2163,9 @@ def fuser_backward( fc1_dy_col_scale = ( fc2_dgrad_kernel_out["sfd_col_tensor"].permute(5, 2, 4, 0, 1, 3).view(-1) ) - grad_scales = fc2_dgrad_kernel_out["dprob_tensor"].view(-1) + grad_scales = fc2_dgrad_kernel_out["dprob_tensor"] + if grad_scales is not None: + grad_scales = grad_scales.view(-1) if recompute_fc2_x_from_dsrelu: d_srelu_tensor = fc2_dgrad_kernel_out.get("d_srelu_tensor") @@ -1856,7 +2188,7 @@ def fuser_backward( fc2_input_quantizer, num_groups, split_sizes, - tensor_offsets=base_split_offsets * fc2_weight_shape[1], + tensor_offsets=fc2_x_tensor_offsets, ) else: sfd_col_d_srelu_tensor = fc2_dgrad_kernel_out.get("sfd_col_d_srelu_tensor") @@ -1878,15 +2210,14 @@ def fuser_backward( scale_inv=None, columnwise_scale_inv=fc2_x_col_scale.reshape(-1), first_dims=split_sizes, - tensor_offsets=base_split_offsets * fc2_weight_shape[1], + tensor_offsets=fc2_x_tensor_offsets, with_gemm_swizzled_scales=True, ) fc2_bias_grads: Optional[list[Optional[torch.Tensor]]] = None fc2_bias_grad_packed: Optional[torch.Tensor] = None if scale_bias: - fc2_biases = fc2_op._get_bias_tensors(dtype) - bias_packed = torch.stack(fc2_biases) + bias_packed = fc2_op._get_packed_bias_tensor(dtype) fc2_dbias_packed_result, grad_scales = compute_grouped_dbias_dscales( fc2_dy, scales_f32, @@ -1921,7 +2252,7 @@ def fuser_backward( fc1_bias_grads = [dbias_2d[group_idx] for group_idx in range(num_groups)] # FC1 grad output for dgrad and wgrad GEMMs - fc1_dy_tensor_offsets = base_split_offsets * fc1_weight_shape[0] + fc1_dy_tensor_offsets = fc1_out_tensor_offsets fc1_grad_output_quantizer = fc1_ctx.grad_output_quantizers[0] if use_nvfp4: fc1_grad_output_quantizer.set_usage( @@ -1992,51 +2323,54 @@ def fuser_backward( if fc1_ctx.input_requires_grad: in_shape = out_shape[:-1] + [fc1_weight_shape[1]] - if use_nvfp4: + fc1_leader = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0 + if is_distributed_weight(fc1_leader): + grouped_fc1_weight = materialize_weight_for_backward(fc1_leader) + + use_single_group_dense_dgrad = num_groups == 1 + if use_single_group_dense_dgrad: grad_input = validate_or_alloc_output(grad_input_buffer, in_shape, dtype, device) - if num_groups == 1: - if fc1_op.single_grouped_weight: - fc1_w_single = grouped_fc1_weight.split_into_quantized_tensors()[0] - else: - fc1_w_single = grouped_fc1_weight[0] - fc1_dy_single = _nvfp4_single_tensor_from_grouped(grouped_fc1_dy) - general_gemm( - fc1_w_single, - fc1_dy_single, - out_dtype=dtype, - out=grad_input, - layout="NN", - ) - else: - fc1_x_tensor_offsets = base_split_offsets * fc1_weight_shape[1] - grouped_grad_input = GroupedTensor( - shape=(out_shape[0], fc1_weight_shape[1]), - dtype=dtype, - num_tensors=num_groups, - quantizer=None, - data=grad_input.view(-1), - first_dims=split_sizes, - tensor_offsets=fc1_x_tensor_offsets, - ) - general_grouped_gemm_for_grouped_tensor( - grouped_fc1_weight, - grouped_fc1_dy, - grouped_grad_input, - layout="NN", - ) + _single_group_dgrad_gemm( + grouped_fc1_dy, + grouped_fc1_weight, + grad_input, + single_grouped_weight=fc1_op.single_grouped_weight, + dtype=dtype, + ) + elif use_nvfp4: + grad_input = validate_or_alloc_output(grad_input_buffer, in_shape, dtype, device) + grouped_grad_input = GroupedTensor( + shape=(out_shape[0], fc1_weight_shape[1]), + dtype=dtype, + num_tensors=num_groups, + quantizer=None, + data=grad_input.view(-1), + first_dims=split_sizes, + tensor_offsets=fc1_x_tensor_offsets, + ) + general_grouped_gemm_for_grouped_tensor( + grouped_fc1_weight, + grouped_fc1_dy, + grouped_grad_input, + layout="NN", + ) else: fc1_dgrad_a_data = fc2_dgrad_kernel_out["d_row_tensor"] fc1_dgrad_a_scales = fc2_dgrad_kernel_out["sfd_row_tensor"] + cudnn_supports_optional_prob = _cudnn_frontend_version_at_least("1.27.0") + fc1_dgrad_prob_tensor = ( + None + if cudnn_supports_optional_prob + else torch.ones((out_shape[0], 1, 1), dtype=torch.float32, device=device) + ) fc1_dgrad_kwargs = { "a_tensor": fc1_dgrad_a_data, "sfa_tensor": fc1_dgrad_a_scales, "padded_offsets": split_points, "alpha_tensor": alpha_tensor, "norm_const_tensor": None, - "prob_tensor": torch.ones( - (out_shape[0], 1, 1), dtype=torch.float32, device=device - ), + "prob_tensor": fc1_dgrad_prob_tensor, "acc_dtype": torch.float32, "d_dtype": dtype, "cd_major": "n", @@ -2045,6 +2379,9 @@ def fuser_backward( "discrete_col_sfd": True, "use_dynamic_sched": True, } + fc1_dgrad_kernel = self.grouped_gemm_quant_kernel() + if _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)): + fc1_dgrad_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 if fc1_op.single_grouped_weight: # Clone and swizzle scales for GEMM @@ -2099,7 +2436,7 @@ def fuser_backward( (out_shape[0], fc1_weight_shape[1], 1), (fc1_weight_shape[1], 1, out_shape[0] * fc1_weight_shape[1]), ) - self.grouped_gemm_quant_kernel()(**fc1_dgrad_kwargs) + fc1_dgrad_kernel(**fc1_dgrad_kwargs) grad_input = grad_input_buffer # FC1 wgrad GEMM @@ -2137,7 +2474,10 @@ def fuser_backward( ) fc2_grad_extra = (None, None) if fc2_op._scale_bias else (None,) - activation_grad_extra = (grad_scales,) if grad_scales is not None else () + if unit_activation_scale: + activation_grad_extra = (None,) + else: + activation_grad_extra = (grad_scales,) if grad_scales is not None else () return ( grad_input, [fc1_grad_params, (), fc2_grad_params], diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py index 42216821c1..1c4b970bd0 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py @@ -5,7 +5,7 @@ """Linear layer forward with Userbuffers communication.""" from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Sequence from typing import Any, Optional import torch @@ -287,7 +287,7 @@ def fuser_forward( prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: # Get basic operations idx = self._op_idxs["linear"] diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 09ffb004dd..fd66529ba8 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -102,23 +102,45 @@ def forward( for tensor in (input_,) + params_and_extra_inputs: tensor._do_not_clear = True - # Unflatten list of parameters and extra tensor inputs - extra_inputs = params_and_extra_inputs[-fuser.num_extra_inputs :] - basic_op_extra_inputs = [] - for op in fuser._basic_ops: - xs, extra_inputs = _split_tuple(extra_inputs, op.num_extra_inputs) - basic_op_extra_inputs.append(xs) + # Place user provided extra inputs into their basic-op slots. Slots bound to + # internal channels are filled lazily as their producers execute. + extra_inputs = params_and_extra_inputs[len(fuser._flat_basic_op_params) :] + basic_op_extra_inputs: list[list[Optional[torch.Tensor]]] = [ + [None] * op.num_extra_inputs for op in fuser._basic_ops + ] + for tensor, (op_idx, input_idx) in zip( + extra_inputs, + fuser._external_extra_input_slots, + ): + basic_op_extra_inputs[op_idx][input_idx] = tensor # Apply forward ops x = input_ - extra_outputs = [None] * fuser._num_basic_ops + extra_outputs: list[Optional[Sequence[Optional[torch.Tensor]]]] = [ + None + ] * fuser._num_basic_ops for op, basic_op_idxs in fuser._forward_ops: # Set if backward op is required for idx in basic_op_idxs: basic_op_ctxs[idx].requires_grad = idx >= fuser.first_op_requiring_backward - # Forward op + # Resolve internal channel inputs from outputs of + # earlier basic ops. When a fusion contains both producer and + # consumer, leave the consumer slot unset so the fused op can + # wire the channel itself + for idx in basic_op_idxs: + for input_idx, source in enumerate(fuser._basic_op_extra_input_sources[idx]): + if source is None: + continue + producer_idx, output_idx = source + if producer_idx in basic_op_idxs: + # fused op will wire the channel itself internally + continue + producer_outputs = extra_outputs[producer_idx] + basic_op_extra_inputs[idx][input_idx] = producer_outputs[output_idx] + + # Prepare args for op forward extra_inputs = [basic_op_extra_inputs[idx] for idx in basic_op_idxs] prev_op_idx = basic_op_idxs[0] - 1 prev_op = fuser._basic_ops[prev_op_idx] if prev_op_idx >= 0 else None @@ -139,24 +161,53 @@ def forward( next_op_input_quantizer=next_op_input_quantizer, basic_op_kwargs=[basic_op_kwargs[idx] for idx in basic_op_idxs], ) + if len(fused_op_extra_outputs) != len(basic_op_idxs): + raise RuntimeError( + f"Expected {type(op).__name__} to generate extra outputs for " + f"{len(basic_op_idxs)} basic operations, " + f"but got {len(fused_op_extra_outputs)}" + ) for idx, ys in zip(basic_op_idxs, fused_op_extra_outputs): - for y in ys: - if set_output_requires_grad: - y.requires_grad_(idx >= fuser.first_op_requiring_backward) + num_extra_outputs = fuser._basic_ops[idx].num_extra_outputs + if len(ys) != num_extra_outputs: + raise RuntimeError( + f"Expected op {idx} to generate {num_extra_outputs} extra outputs, " + f"but got {len(ys)}" + ) + for output_idx, y in enumerate(ys): + if y is None: + # Extra output can be None if it is not required by any operations outside the fusion + # and is not required to be outputted to the caller. + output_to_caller = fuser._basic_op_extra_output_to_caller[idx][output_idx] + consumers = fuser._basic_op_extra_output_consumers[idx][output_idx] + needed_outside_fusion = any( + consumer_idx not in basic_op_idxs for consumer_idx in consumers + ) + if output_to_caller: + raise RuntimeError( + f"Op {idx} extra output {output_idx} is public, " + f"but {type(op).__name__} returned None" + ) + if needed_outside_fusion: + raise RuntimeError( + f"Op {idx} extra output {output_idx} is required by an " + "operation outside its forward fusion, " + f"but {type(op).__name__} returned None" + ) + continue + if ( + set_output_requires_grad + and idx >= fuser.first_op_requiring_backward + and y.is_floating_point() + ): + y.requires_grad_(True) extra_outputs[idx] = ys - # Flatten list of extra outputs - extra_outputs_flat = [] - for idx, ys in enumerate(extra_outputs): - ys = list(ys) - num_extra_outputs = fuser._basic_ops[idx].num_extra_outputs - if len(ys) != num_extra_outputs: - raise RuntimeError( - f"Expected op {idx} to generate " - "{num_extra_outputs} extra inputs, " - f"but got {len(ys)}" - ) - extra_outputs_flat.extend(ys) + # Collect caller-visible extra outputs in basic-op and slot order. + extra_outputs_flat = [ + extra_outputs[op_idx][output_idx] + for op_idx, output_idx in fuser._public_extra_output_slots + ] # Save context for backward pass if func_ctx is not None: @@ -186,12 +237,19 @@ def forward( func_ctx.basic_ops = fuser._basic_ops func_ctx.basic_op_ctxs = basic_op_ctxs func_ctx.basic_op_num_params = fuser._basic_op_num_params - func_ctx.num_extra_inputs = fuser.num_extra_inputs func_ctx.num_extra_outputs = len(extra_outputs_flat) + func_ctx.external_extra_input_slots = fuser._external_extra_input_slots + func_ctx.public_extra_output_slots = fuser._public_extra_output_slots + func_ctx.basic_op_extra_output_channels = fuser._basic_op_extra_output_channels + func_ctx.basic_op_extra_output_consumers = fuser._basic_op_extra_output_consumers + func_ctx.basic_op_extra_input_sources = fuser._basic_op_extra_input_sources func_ctx.is_first_module = is_first_module # Mark output tensors as not deletable in backward - for tensor in [x] + extra_outputs_flat: + for tensor in itertools.chain( + (x,), + (y for ys in extra_outputs for y in ys if y is not None), + ): tensor._do_not_clear = True if set_output_requires_grad: @@ -224,21 +282,32 @@ def backward( ctx.saved_tensors = saved_tensors[slice(*ctx._saved_tensors_range)] ctx._saved_tensors_range = None - # Unflatten list of extra tensor output grads + # Channel wiring saved from forward + basic_op_extra_output_channels = func_ctx.basic_op_extra_output_channels + basic_op_extra_output_consumers = func_ctx.basic_op_extra_output_consumers + basic_op_extra_input_sources = func_ctx.basic_op_extra_input_sources + + # Place caller-provided extra-output grads into their basic-op slots. + # Gradients from internal channel consumers are added during backward. if len(grad_extra_outputs) != func_ctx.num_extra_outputs: raise ValueError( f"Expected grads for {func_ctx.num_extra_outputs} extra tensor outputs, " f"but got {len(grad_extra_outputs)}" ) - basic_op_grad_extra_outputs = [] - for op in basic_ops: - dys, grad_extra_outputs = _split_tuple(grad_extra_outputs, op.num_extra_outputs) - basic_op_grad_extra_outputs.append(dys) + basic_op_grad_extra_outputs: list[list[Optional[torch.Tensor]]] = [ + [None] * op.num_extra_outputs for op in basic_ops + ] + for grad, (op_idx, output_idx) in zip( + grad_extra_outputs, + func_ctx.public_extra_output_slots, + ): + basic_op_grad_extra_outputs[op_idx][output_idx] = grad # Apply backward ops dx = grad_output grad_params = [None for _ in range(len(basic_ops))] grad_extra_inputs = [None for _ in range(len(basic_ops))] + channel_grads: dict[str, torch.Tensor] = {} for op, basic_op_idxs in reversed(backward_ops): # Stop if no more gradients are required @@ -246,7 +315,17 @@ def backward( dx = None break - # Backward op + # Backward op. Supply gradients accumulated from every consumer of + # each internal channel. + for idx in basic_op_idxs: + for output_idx, channel in enumerate(basic_op_extra_output_channels[idx]): + if basic_op_extra_output_consumers[idx][output_idx]: + channel_grad = channel_grads.get(channel) + if channel_grad is not None: + output_grad = basic_op_grad_extra_outputs[idx][output_idx] + basic_op_grad_extra_outputs[idx][output_idx] = ( + channel_grad if output_grad is None else output_grad + channel_grad + ) grad_extra_outputs = [basic_op_grad_extra_outputs[idx] for idx in basic_op_idxs] dx, fused_op_grad_params, fused_op_grad_extra_inputs = op.fuser_backward( [basic_op_ctxs[idx] for idx in basic_op_idxs], @@ -258,6 +337,18 @@ def backward( basic_op_ctxs[idx].saved_tensors = None for idx, dxs in zip(basic_op_idxs, fused_op_grad_extra_inputs): grad_extra_inputs[idx] = dxs + for input_idx, grad in enumerate(dxs): + source = basic_op_extra_input_sources[idx][input_idx] + if source is None or grad is None: + continue + producer_idx, output_idx = source + # Producer already ran inside this fusion; the fused op + # must apply these grads itself rather than via channel_grads. + if producer_idx in basic_op_idxs: + continue + channel = basic_op_extra_output_channels[producer_idx][output_idx] + previous_grad = channel_grads.get(channel) + channel_grads[channel] = grad if previous_grad is None else previous_grad + grad # Flatten list of parameter gradients grad_params_flat = [] @@ -275,20 +366,22 @@ def backward( grad_params_flat.extend(dparams) # Flatten list of parameter gradients - grad_extra_inputs_flat = [] for idx, dxs in enumerate(grad_extra_inputs): num_extra_inputs = basic_ops[idx].num_extra_inputs if dxs is None: - dxs = [None for _ in range(num_extra_inputs)] - else: - dxs = list(dxs) - if len(dxs) != num_extra_inputs: + grad_extra_inputs[idx] = (None,) * num_extra_inputs + elif len(dxs) != num_extra_inputs: raise RuntimeError( f"Expected op {idx} to generate grads " f"for {num_extra_inputs} extra inputs, " f"but got {len(dxs)}" ) - grad_extra_inputs_flat.extend(dxs) + + # Collect the gradient for each public extra input. + grad_extra_inputs_flat = [ + grad_extra_inputs[op_idx][input_idx] + for op_idx, input_idx in func_ctx.external_extra_input_slots + ] # Update FP8 scaling factors if func_ctx.is_first_module and not _is_graph_capturing(): @@ -342,7 +435,104 @@ def __init__( # Number of extra tensor inputs self._basic_op_num_extra_inputs: list[int] = list(op.num_extra_inputs for op in basic_ops) - self.num_extra_inputs: int = sum(self._basic_op_num_extra_inputs) + self._basic_op_extra_input_sources: list[list[Optional[tuple[int, int]]]] = [ + [None] * op.num_extra_inputs for op in basic_ops + ] + self._basic_op_extra_output_channels: list[list[Optional[str]]] = [ + list(op._extra_output_channels) for op in basic_ops + ] + self._basic_op_extra_output_to_caller: list[list[bool]] = [ + list(op._extra_output_to_caller) for op in basic_ops + ] + self._basic_op_extra_output_consumers: list[list[list[int]]] = [ + [[] for _ in range(op.num_extra_outputs)] for op in basic_ops + ] + self._external_extra_input_slots: list[tuple[int, int]] = [] + self._public_extra_output_slots: list[tuple[int, int]] = [] + + # Find channel producers and reject ambiguous names. + channel_producers: dict[str, tuple[int, int]] = {} + for op_idx, op in enumerate(basic_ops): + for output_idx, channel in enumerate(self._basic_op_extra_output_channels[op_idx]): + if channel is None: + continue + if channel in channel_producers: + producer_idx, _ = channel_producers[channel] + raise ValueError( + f"Extra tensor channel {channel!r} has multiple producers " + f"(ops {producer_idx} and {op_idx})" + ) + channel_producers[channel] = (op_idx, output_idx) + + # Resolve inputs. Named inputs must have an earlier producer; + # unnamed inputs remain public. + for op_idx, op in enumerate(basic_ops): + for input_idx, channel in enumerate(op._extra_input_channels): + if channel is None: + self._external_extra_input_slots.append((op_idx, input_idx)) + continue + producer = channel_producers.get(channel) + if producer is None: + raise ValueError( + f"Extra tensor channel {channel!r} consumed by op {op_idx} " + f"({type(op).__name__}) has no producer" + ) + producer_idx, _ = producer + if producer_idx >= op_idx: + raise ValueError( + f"Extra tensor channel {channel!r} consumed by op {op_idx} " + f"({type(op).__name__}) has no earlier producer" + ) + self._basic_op_extra_input_sources[op_idx][input_idx] = producer + producer_idx, output_idx = producer + self._basic_op_extra_output_consumers[producer_idx][output_idx].append(op_idx) + + # Record caller-visible outputs in stable basic-op and slot order. + for op_idx, op in enumerate(basic_ops): + for output_idx in range(op.num_extra_outputs): + if self._basic_op_extra_output_to_caller[op_idx][output_idx]: + self._public_extra_output_slots.append((op_idx, output_idx)) + + # Every channel-bound extra input must be wired to a matching producer + # extra output. External slots remain unbound (source is None). + for op_idx, sources in enumerate(self._basic_op_extra_input_sources): + op = basic_ops[op_idx] + for input_idx, source in enumerate(sources): + channel = op._extra_input_channels[input_idx] + if channel is None: + if source is not None: + raise RuntimeError( + f"Extra input {input_idx} of op {op_idx} " + f"({type(op).__name__}) is external but has a " + f"producer source {source}" + ) + continue + if source is None: + raise RuntimeError( + f"Extra input {input_idx} of op {op_idx} " + f"({type(op).__name__}) is bound to channel {channel!r} " + "without a producer source" + ) + producer_idx, output_idx = source + producer_channel = self._basic_op_extra_output_channels[producer_idx][output_idx] + if producer_channel != channel: + raise ValueError( + f"Extra input {input_idx} of op {op_idx} " + f"({type(op).__name__}) is bound to channel {channel!r}, " + f"but producer op {producer_idx} extra output {output_idx} " + f"is bound to {producer_channel!r}" + ) + # Used by Sequential to determine the number of extra inputs + # needed for each OperationFuser module in the sequence. + self.num_extra_inputs = len(self._external_extra_input_slots) + + # This fuser's routing is derived from the channel bindings above, so + # freeze them for every covered basic op. That includes transient + # fusers from BasicOperation.forward / FusedOperation.forward + # (op(x)), not only persistent OperationFuser / Sequential usage. + # Changing bindings afterward requires constructing new operations. + for op in self._basic_ops: + op._lock_extra_tensor_channels() # Ops for forward and backward pass, will be populated in maybe_fuse_ops self._forward_ops: list[tuple[FusibleOperation, list[int]]] @@ -432,7 +622,7 @@ def maybe_fuse_ops( first_op_requiring_backward = self._num_basic_ops for op_idx in range(self._num_basic_ops): op_inputs = itertools.chain(self._basic_op_params[op_idx], extra_inputs[op_idx]) - if any(tensor.requires_grad for tensor in op_inputs): + if any(tensor is not None and tensor.requires_grad for tensor in op_inputs): first_op_requiring_backward = op_idx break @@ -507,6 +697,7 @@ def __call__( *extra_inputs: torch.Tensor, basic_op_kwargs: Optional[list[dict[str, Any]]] = None, ) -> torch.Tensor | tuple[torch.Tensor, ...]: + # Verify extra input count if len(extra_inputs) != self.num_extra_inputs: raise ValueError( @@ -517,12 +708,13 @@ def __call__( if basic_op_kwargs is None: basic_op_kwargs = [{}] * self._num_basic_ops - # Unflatten list of extra tensor inputs - extra_inputs_copy = list(extra_inputs) - basic_op_extra_inputs = [] - for op in self._basic_ops: - xs, extra_inputs_copy = _split_tuple(extra_inputs_copy, op.num_extra_inputs) - basic_op_extra_inputs.append(xs) + # Place public extra inputs into their basic-op slots. Internal slots + # are not available until forward executes their producers. + basic_op_extra_inputs: list[list[Optional[torch.Tensor]]] = [ + [None] * op.num_extra_inputs for op in self._basic_ops + ] + for tensor, (op_idx, input_idx) in zip(extra_inputs, self._external_extra_input_slots): + basic_op_extra_inputs[op_idx][input_idx] = tensor # Get environment state recipe = None diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 849c900f95..d057d46816 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -6,7 +6,7 @@ from __future__ import annotations import abc -from collections.abc import Iterable +from collections.abc import Iterable, Sequence import dataclasses import pickle from typing import Any, Optional @@ -85,11 +85,11 @@ def fuser_forward( basic_op_ctxs: list[OperationContext], input_: torch.Tensor, *, - basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], + basic_op_extra_inputs: Sequence[Sequence[Optional[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, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[Optional[torch.Tensor]]]]: """Forward pass This op is either a basic op or the fusion of basic ops, so @@ -104,8 +104,9 @@ def fuser_forward( Contexts for basic operations input_: torch.Tensor Input tensor - basic_op_extra_inputs: list of torch.Tensor - Extra tensor inputs to basic operations + basic_op_extra_inputs: sequence of sequences of torch.Tensor + Extra tensor inputs to basic operations. An internal input + owned by this fused operation may be ``None``. prev_op_grad_output_quantizer: Quantizer, optional The grad_output_quantizer of the preceeding operation next_op_input_quantizer: Quantizer, optional @@ -118,8 +119,9 @@ def fuser_forward( ------- torch.Tensor: Output tensor. - Iterable of torch.Tensor: - Extra tensor outputs from basic operations. + Sequence of sequences of torch.Tensor: + Extra tensor outputs from basic operations. A non-public + channel owned by this fused operation may be ``None``. """ raise NotImplementedError( @@ -131,11 +133,11 @@ def fuser_backward( basic_op_ctxs: list[OperationContext], grad_output: torch.Tensor, *, - basic_op_grad_extra_outputs: list[tuple[torch.Tensor, ...]], + basic_op_grad_extra_outputs: Sequence[Sequence[Optional[torch.Tensor]]], ) -> tuple[ torch.Tensor, - Iterable[Iterable[Optional[torch.Tensor]]], - Iterable[Iterable[Optional[torch.Tensor]]], + Sequence[Sequence[Optional[torch.Tensor]]], + Sequence[Sequence[Optional[torch.Tensor]]], ]: """Backward pass @@ -159,9 +161,9 @@ def fuser_backward( ------- torch.Tensor: Loss gradient w.r.t. operation input - Iterable of iterable of torch.Tensor: + Sequence of sequences of torch.Tensor: Loss gradients w.r.t. parameters for basic operations - Iterable of iterable of torch.Tensor: + Sequence of sequences of torch.Tensor: Loss gradients w.r.t. extra tensor inputs to basic operations @@ -187,10 +189,91 @@ class BasicOperation(FusibleOperation, metaclass=abc.ABCMeta): def __init__(self) -> None: super().__init__() + # Optional names for extra-tensor channels internal to an OperationFuser. + # Unbound slots remain public inputs/outputs, preserving the original API. + self._extra_input_channels: list[Optional[str]] = [None] * self.num_extra_inputs + self._extra_output_channels: list[Optional[str]] = [None] * self.num_extra_outputs + self._extra_output_to_caller: list[bool] = [True] * self.num_extra_outputs + # Channel routing is captured by an OperationFuser when it is + # constructed (including the transient fuser created by a + # standalone op(x) call), so it is frozen once that happens. + self._extra_tensor_channels_locked: bool = False + # Objects for quantization self._fp8_metas: Optional[dict[str, dict[str, Any]]] = None self._quantizers: Optional[dict[str, list[Quantizer]]] = None + def _lock_extra_tensor_channels(self) -> None: + """Freeze channel routing after an OperationFuser has captured it.""" + self._extra_tensor_channels_locked = True + + def _check_extra_tensor_channels_unlocked(self) -> None: + """Reject channel rebinding after an OperationFuser has captured it.""" + if self._extra_tensor_channels_locked: + raise RuntimeError( + f"Cannot change extra tensor channels of {type(self).__name__} because an " + "OperationFuser has already captured its channel routing. Construct new " + "operations, bind their channels, and build a new OperationFuser or Sequential." + ) + + def set_extra_input_channel(self, index: int, channel: Optional[str]) -> BasicOperation: + """Bind an extra input slot to an internal fuser channel. + + A bound slot receives the matching extra output from an earlier + operation in the same fuser instead of consuming a public extra input. + Passing ``None`` removes the binding. + + Channels must be bound before any ``OperationFuser`` captures + them. That includes constructing an ``OperationFuser`` or + ``Sequential``, and also calling the op directly (``op(x)``), + which builds a transient fuser and locks channels permanently. + """ + if not 0 <= index < self.num_extra_inputs: + raise IndexError( + f"Extra input index {index} is out of range for " + f"{type(self).__name__} with {self.num_extra_inputs} extra inputs" + ) + if channel is not None and (not isinstance(channel, str) or not channel): + raise ValueError("Extra input channel must be a non-empty string or None") + self._check_extra_tensor_channels_unlocked() + self._extra_input_channels[index] = channel + return self + + def set_extra_output_channel( + self, + index: int, + channel: Optional[str], + *, + output_to_caller: bool = True, + ) -> BasicOperation: + """Bind an extra output slot to an internal fuser channel. + + A bound slot can feed one or more later operations. By default, the + output is also returned to the caller. Set ``output_to_caller=False`` + to keep it internal to the fuser. Passing ``channel=None`` removes the + binding and restores the output as public. + + Channels must be bound before any ``OperationFuser`` captures + them. That includes constructing an ``OperationFuser`` or + ``Sequential``, and also calling the op directly (``op(x)``), + which builds a transient fuser and locks channels permanently. + """ + if not 0 <= index < self.num_extra_outputs: + raise IndexError( + f"Extra output index {index} is out of range for " + f"{type(self).__name__} with {self.num_extra_outputs} extra outputs" + ) + if channel is not None and (not isinstance(channel, str) or not channel): + raise ValueError("Extra output channel must be a non-empty string or None") + if not isinstance(output_to_caller, bool): + raise TypeError("output_to_caller must be a bool") + if channel is None: + output_to_caller = True + self._check_extra_tensor_channels_unlocked() + self._extra_output_channels[index] = channel + self._extra_output_to_caller[index] = output_to_caller + return self + @property def is_fused_op(self) -> bool: return False @@ -275,11 +358,6 @@ def reset_recipe_state( if num_quantizers == 0: continue - if recipe.float8_block_scaling(): - raise NotImplementedError( - "Fusible operations do not support FP8 block scaling recipe" - ) - # Construct quantization recipe state roles = self.get_quantizer_roles(mode) # pylint: disable=assignment-from-none if roles is not None: @@ -537,7 +615,12 @@ def forward( *extra_inputs: torch.Tensor, **kwargs: Any, ) -> torch.Tensor | tuple[torch.Tensor, ...]: - """Apply operation""" + """Apply operation. + + Builds a transient ``OperationFuser([self])``, which captures and + locks this op's extra-tensor channel bindings. Bind channels before + the first call if they will be used later in a multi-op fuser. + """ from .fuser import OperationFuser return OperationFuser([self])( @@ -770,7 +853,11 @@ def forward( *extra_inputs: torch.Tensor, basic_op_kwargs: Optional[list[dict[str, Any]]] = None, ) -> torch.Tensor: - """Apply operation""" + """Apply operation. + + Builds a transient ``OperationFuser([self])``, which captures and + locks every basic op's extra-tensor channel bindings. + """ if basic_op_kwargs is None: basic_op_kwargs = [{} for _ in range(len(self.basic_ops))] from .fuser import OperationFuser diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py index 2d482576ad..ef65081a73 100644 --- a/transformer_engine/pytorch/optimizers/fused_adam.py +++ b/transformer_engine/pytorch/optimizers/fused_adam.py @@ -571,12 +571,11 @@ def step(self, closure=None, grad_scaler=None): loss = closure() for group in self.param_groups: - if len(group["params"]) == 0: - continue - device = group["params"][0].device - bias_correction = 1 if group["bias_correction"] else 0 - beta1, beta2 = group["betas"] - + # Advance the step counter before skipping empty groups. A param group can be + # empty on some data-parallel ranks and populated on others, so incrementing + # only for populated groups desynchronizes "step" across the ranks that share + # an optimizer state shard. A rank then resumes from a checkpoint written by a + # rank where the group was empty and applies a stale bias correction. # assume same step across group now to simplify things # per parameter step can be easily support by making it tensor, or pass list into kernel if "step" in group: @@ -584,10 +583,25 @@ def step(self, closure=None, grad_scaler=None): 1 if not self.capturable else (self._dummy_overflow_buf != 1).to(torch.int) ) else: + # Empty groups have no parameter to take the device from, so fall back to + # the device of the optimizer's own scratch buffer. + step_device = ( + group["params"][0].device + if len(group["params"]) > 0 + else self._dummy_overflow_buf.device + ) group["step"] = ( - 1 if not self.capturable else torch.tensor([1], dtype=torch.int, device=device) + 1 + if not self.capturable + else torch.tensor([1], dtype=torch.int, device=step_device) ) + if len(group["params"]) == 0: + continue + device = group["params"][0].device + bias_correction = 1 if group["bias_correction"] else 0 + beta1, beta2 = group["betas"] + # create lists for multi-tensor apply p_main_of_fp8_model = [] p_main_of_f16_model = [] diff --git a/transformer_engine/pytorch/optimizers/newton_schulz.py b/transformer_engine/pytorch/optimizers/newton_schulz.py new file mode 100644 index 0000000000..4f868f1c12 --- /dev/null +++ b/transformer_engine/pytorch/optimizers/newton_schulz.py @@ -0,0 +1,353 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Distributed Newton-Schulz matrix orthogonalization via cuSolverMp.""" + +from itertools import chain, cycle, islice, repeat +from typing import Iterator, Literal, Optional, Sequence + +import torch +import torch.distributed as dist + +import transformer_engine_torch as tex + + +_COEFFICIENT_SETS = { + # Values are rounded to closest representable in single precision. + "simple": [ + (3.4445, -4.7750, 2.0315), + ], + "quintic": [ + # optimized for a quintic iteration. + # Source: https://leloykun.github.io/ponder/muon-opt-coeffs/#how-do-we-optimize-the-coefficients + # Numbers from: https://github.com/KellerJordan/modded-nanogpt/blob/0674386070ceb4dcd207e1aca747ffcea6c15250/train_gpt_medium.py#L45 + (4.0848, -6.8946, 2.9270), + (3.9505, -6.3029, 2.6377), + (3.7418, -5.5913, 2.3037), + (2.8769, -3.1427, 1.2046), + (2.8366, -3.0525, 1.2012), + ], + "polar_express": [ + # Polar Express iteration from: https://arxiv.org/abs/2505.16932 + # We include PolarExpress' division by 1.01^polynomial_degree (as stated in their Algorithm 1) in the coefficient list. + # This is a safety factor for numerical stability. + (8.2051, -22.9019, 16.4607), + (4.0664, -2.8612, 0.5184), + (3.9096, -2.8234, 0.5250), + (3.2856, -2.4153, 0.4853), + (2.2779, -1.6198, 0.3985), + (1.8726, -1.2307, 0.3585), + (1.8564, -1.2132, 0.3568), + (1.8750, -1.2500, 0.3750), + ], + "cans": [ + # CANS from: http://arxiv.org/abs/2506.10935 + # CANS iteration (Remez + adaptive interval) based coefficients. + # Source (for generating CANS coefficients): https://github.com/GrishKate/accelerating_orthogonalization/blob/main/polynomials.py + (8.4703, -25.1081, 18.6293), + (4.1828, -3.1087, 0.5806), + (3.9619, -2.9541, 0.5630), + (3.2866, -2.4647, 0.5074), + (2.2737, -1.6447, 0.4162), + ], + "aol": [ + # from https://github.com/thib-s/flash-newton-schulz/blob/main/newton_schulz_triton.py#L511 + (4.0098, -7.0585, 2.4635), + (3.4585, -5.5479, 2.5959), + (2.7573, -3.2939, 1.4254), + (2.7215, -3.0494, 1.3169), + ], +} + +NSCoeffT = Literal[_COEFFICIENT_SETS.keys()] + +CoeffIterMode = Literal["cycle", "repeat_last"] +CoeffT = tuple[float, float, float] + + +def get_coefficient_iterator( + steps: int, + coefficient_sets: Sequence[CoeffT], + mode: CoeffIterMode = "cycle", +) -> Iterator[CoeffT]: + """Iterate through coefficient sets with configurable end behavior using itertools. + + Args: + steps: The number of tuples to yield. + coefficient_sets: A sequence of (a, b, c) coefficient tuples. + mode: Iteration mode: + - "cycle": After the last element, restart from the beginning. + - "repeat_last": After the last element, keep yielding the last tuple. + + Yields: + Tuples (a, b, c) from coefficient_sets according to the specified mode. + + Raises: + ValueError: If coefficient_sets is empty. + ValueError: If an invalid mode is provided. + """ + if not coefficient_sets: + raise ValueError("coefficient_sets must be non-empty.") + + base: Iterator[CoeffT] + if mode == "cycle": + base = cycle(coefficient_sets) + elif mode == "repeat_last": + # Chain the original list with an infinite repeat of the last item + base = chain(coefficient_sets, repeat(coefficient_sets[-1])) + else: + raise ValueError(f"Invalid mode: {mode}. Expected 'cycle' or 'repeat_last'.") + + return islice(base, steps) + + +def get_coefficients(steps: int, coefficient_type: NSCoeffT = "quintic") -> list[CoeffT]: + """Return the coefficient schedule for Newton-Schulz. + + Parameter ``coefficient_type`` can be one of the following + - "simple": Default coefficient set. + - "quintic": Quintic iteration with optimized coefficients. + - "polar_express": Polar Express iteration with optimized coefficients. + - "cans": CANS iteration with Remez + adaptive interval coefficients. + - "aol": AOL coefficient set. + """ + if coefficient_type not in _COEFFICIENT_SETS: + raise ValueError("Invalid coefficient type: " + coefficient_type) + iter_mode: CoeffIterMode = ( + "repeat_last" if coefficient_type in ("polar_express", "cans") else "cycle" + ) + coeff_iter = get_coefficient_iterator( + steps, _COEFFICIENT_SETS[coefficient_type], mode=iter_mode + ) + return list(coeff_iter) + + +class CusolverMpCtx: + """cuSolverMp context for Newton-Schulz matrix orthogonalization. + + Context creation is expensive; create once and reuse across multiple + :func:`newton_schulz` calls. Creation is collective over ``group`` and + must be called by every group member on its intended CUDA device. Call + :meth:`destroy` before destroying ``group``. + """ + + def __init__(self, group: dist.ProcessGroup) -> None: + # The cuSolverMp grid borrows the ProcessGroupNCCL communicator. Keep + # the group alive until the native context has released the grid. + self._ptr: Optional[int] = None + self._group: Optional[dist.ProcessGroup] = group + + if not dist.is_initialized(): + raise RuntimeError( + "torch.distributed must be initialized before creating CusolverMpCtx" + ) + + self.nranks = dist.get_world_size(group) + self.rank = dist.get_rank(group) + if self.rank < 0: + raise RuntimeError("The current process is not a member of the supplied process group") + + comm_ptr = _get_nccl_comm_ptr(group) + self._ptr = tex.cusolvermp_ctx_create(comm_ptr, self.nranks, self.rank) + + @property + def group(self) -> dist.ProcessGroup: + """Process group backing the cuSolverMp context.""" + if self._group is None: + raise RuntimeError("CusolverMpCtx has been destroyed") + return self._group + + def destroy(self) -> None: + """Destroy the underlying cuSolverMp context.""" + try: + if self._ptr is not None: + tex.cusolvermp_ctx_destroy(self._ptr) + finally: + self._ptr = None + self._group = None + + def __del__(self) -> None: + # Called when the context is manually destroyed or during Python teardown + self.destroy() + + +def _get_nccl_comm_ptr(group: dist.ProcessGroup) -> int: + """Materialize and borrow a raw NCCL communicator from a process group.""" + backend = dist.get_backend(group) + if backend != "nccl": + raise RuntimeError(f"Newton-Schulz requires NCCL backend, got '{backend}'") + + # ProcessGroupNCCL creates communicators lazily. This device-specific + # collective ensures that the communicator returned by _comm_ptr() exists + # and is ready on every group rank before cuSolverMp borrows it. + dist.barrier(group=group, device_ids=[torch.cuda.current_device()]) + nccl_backend = group._get_backend(torch.device("cuda")) + comm_ptr = nccl_backend._comm_ptr() + if not isinstance(comm_ptr, int) or comm_ptr == 0: + raise RuntimeError("ProcessGroupNCCL returned an invalid communicator pointer") + return comm_ptr + + +def newton_schulz( + x: torch.Tensor, + ctx: CusolverMpCtx, + num_iterations: int = 5, + coefficients: Optional[Sequence[CoeffT]] = None, +) -> None: + """Compute Newton-Schulz matrix orthogonalization in-place on a distributed matrix. + + Parameters + ---------- + x : torch.Tensor + Local part of the distributed matrix (modified in-place). + Must be a 2D CUDA tensor of type float32 or bfloat16. + Columns are distributed across ranks. + ctx : CusolverMpCtx + cuSolverMp context created by :func:`cusolvermp_ctx_create`. + num_iterations : int, optional + Number of Newton-Schulz iterations. Default: 5. + coefficients : sequence of tuple[float, float, float], optional + Polynomial coefficients for the Newton-Schulz iteration. + """ + if coefficients is None: + coefficients = get_coefficients(num_iterations) + if len(coefficients) != num_iterations: + raise ValueError( + f"Unexpected number of coefficients: {len(coefficients)} for" + f" {num_iterations} iterations" + ) + flat_coefficients: list[float] = [] + for i, coeff in enumerate(coefficients): + if len(coeff) != 3: + raise ValueError( + f"Expected coefficient tuple of length 3 at iteration {i}, got {len(coeff)}" + ) + flat_coefficients.extend(coeff) + + if x.dim() != 2: + raise ValueError(f"Expected 2D tensor, got {x.dim()}D") + if x.dtype not in (torch.float32, torch.bfloat16): + raise ValueError(f"Expected float32 or bfloat16 tensor, got {x.dtype}") + if not x.is_contiguous(): + raise ValueError("Input tensor must be contiguous") + if not x.is_cuda: + raise ValueError("Input tensor must be on CUDA device") + + # Global matrix dimensions; columns are distributed across ranks. + m = x.size(0) + n = x.size(1) * ctx.nranks + + tex.newton_schulz(ctx._ptr, m, n, x, num_iterations, flat_coefficients) + + +def _orthogonalize_replicated( + x: torch.Tensor, + ctx: CusolverMpCtx, + num_iterations: int, + coefficients: Optional[Sequence[CoeffT]], + transpose: bool, +) -> torch.Tensor: + """Orthogonalize a replicated matrix using the distributed column-sharded kernel.""" + work = x.mT.contiguous() if transpose else x + if work.size(1) % ctx.nranks != 0: + distributed_dim = 0 if transpose else 1 + raise ValueError( + f"Tensor dimension {distributed_dim} with size {x.size(distributed_dim)} " + f"must be divisible by tensor-parallel size {ctx.nranks}" + ) + + local_work = work.chunk(ctx.nranks, dim=1)[ctx.rank].contiguous() + newton_schulz(local_work, ctx, num_iterations, coefficients=coefficients) + + output_shards = [torch.empty_like(local_work) for _ in range(ctx.nranks)] + dist.all_gather(output_shards, local_work, group=ctx.group) + output = torch.cat(output_shards, dim=1) + return output.mT.contiguous() if transpose else output + + +def newton_schulz_tp( + x: torch.Tensor, + ctx: CusolverMpCtx, + num_iterations: int = 5, + coefficients: Optional[Sequence[CoeffT]] = None, + partition_dim: Optional[int] = None, + tp_mode: Literal["duplicated", "distributed"] = "duplicated", +) -> None: + """Compute tensor-parallel Newton-Schulz orthogonalization in-place. + + This convenience wrapper handles replicated tensors and tensor-parallel + shards while delegating the matrix orthogonalization to :func:`newton_schulz`. + The underlying kernel expects columns to be distributed across ranks, so row + partitions are transposed before the call and transposed back afterward. + + Parameters + ---------- + x : torch.Tensor + Local tensor to orthogonalize in-place. + ctx : CusolverMpCtx + cuSolverMp context created for the tensor-parallel process group. + num_iterations : int, optional + Number of Newton-Schulz iterations. Default: 5. + coefficients : sequence of tuple[float, float, float], optional + Polynomial coefficients for the Newton-Schulz iteration. + partition_dim : int, optional + Dimension along which ``x`` is partitioned. ``None`` treats ``x`` as a + duplicated full tensor on every rank. + tp_mode : {"duplicated", "distributed"}, optional + ``"distributed"`` orthogonalizes the existing partition directly. + ``"duplicated"`` first gathers the full tensor, orthogonalizes it, and + copies this rank's partition back into ``x``. + """ + if x.dim() != 2: + raise ValueError(f"Expected 2D tensor, got {x.dim()}D") + + if partition_dim is not None: + if partition_dim not in (0, 1): + raise ValueError(f"Invalid partition_dim: {partition_dim}") + if tp_mode not in ("duplicated", "distributed"): + raise ValueError(f"Invalid tp_mode: {tp_mode}") + + if x.dtype not in (torch.float32, torch.bfloat16): + raise ValueError(f"Expected float32 or bfloat16 tensor, got {x.dtype}") + if not x.is_contiguous(): + raise ValueError("Input tensor must be contiguous") + if not x.is_cuda: + raise ValueError("Input tensor must be on CUDA device") + + if partition_dim is None: + output = _orthogonalize_replicated( + x, + ctx, + num_iterations, + coefficients, + transpose=x.size(0) > x.size(1), + ) + x.copy_(output) + return + + if tp_mode == "duplicated": + x_shards = [torch.empty_like(x) for _ in range(ctx.nranks)] + dist.all_gather(x_shards, x, group=ctx.group) + global_x = torch.cat(x_shards, dim=partition_dim) + + output = _orthogonalize_replicated( + global_x, + ctx, + num_iterations, + coefficients, + transpose=global_x.size(0) > global_x.size(1), + ) + + local_start = ctx.rank * x.size(partition_dim) + local_output = output.narrow(partition_dim, local_start, x.size(partition_dim)) + x.copy_(local_output) + elif tp_mode == "distributed": + if partition_dim == 0: + x_t = x.mT.contiguous() + newton_schulz(x_t, ctx, num_iterations, coefficients=coefficients) + x.copy_(x_t.mT) + else: + newton_schulz(x, ctx, num_iterations, coefficients=coefficients) + else: + raise ValueError(f"Invalid tp_mode: {tp_mode}") diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index b11357a766..c3802a08fb 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -324,7 +324,17 @@ def get_default_recipe() -> Recipe: def get_align_size_for_quantization(recipe: Recipe) -> int: - """Get the alignment size for quantization.""" + """Get the alignment used to pad grouped quantized operations. + + Built-in recipes use their format requirement. Custom recipes use their + declarative ``quantization_alignment`` contract rather than invoking + ``qfactory``, since a factory may be stateful or role-dependent. + """ + # TODO(#3158): Prefer module/role-specific alignment derived from canonical + # cached quantizers when that context is available. Keep the recipe-wide + # alignment as the conservative fallback for context-free callers. + if recipe.custom(): + return recipe.quantization_alignment if recipe.mxfp8(): # HipKittens grouped GEMM requires 256-aligned expert dimensions. # HK is used by default when NVTE_USE_CUTLASS_GROUPED_GEMM=1, @@ -1644,7 +1654,23 @@ def make_quantizers(self) -> list: # TODO(ksivamani); Find better design for this, adding here to avoid circular import. from .tensor.mxfp8_tensor import MXFP8Quantizer - return [MXFP8Quantizer(self.dtype) for i in range(self.num_quantizers)] + if self.mode not in ("forward", "backward"): + raise RuntimeError(f"Unexpected recipe mode ({self.mode})") + + if self.mode == "backward" or not self.recipe.enable_2d_quantization: + return [MXFP8Quantizer(self.dtype) for i in range(self.num_quantizers)] + + def _use_2d_quantization(idx: int) -> bool: + role = self._slot_role(idx) + return role.module_type in ("linear", "grouped_linear") and role.tensor_type == "weight" + + return [ + MXFP8Quantizer( + self.dtype, + with_2d_quantization=_use_2d_quantization(idx), + ) + for idx in range(self.num_quantizers) + ] class Float8BlockScalingRecipeState(RecipeState): @@ -2097,18 +2123,17 @@ def make_quantizers(self) -> list: ) roles = [QuantizerRole() for _ in range(self.num_quantizers)] - # qfactory must return a Quantizer or QuantizerRequest for every slot. - # None is not a valid return value — it would silently disable quantization - # for that tensor, risking hard-to-detect performance regressions. - # TODO(negvet): Introduce an explicit IdentityQuantizer for intentional no-op - # quantization. Until then, None is rejected. + # qfactory returns one quantizer-like object per slot; use + # ``IdentityQuantizer`` for intentional high-precision passthrough. raw = [qfactory(roles[i]) for i in range(self.num_quantizers)] for i, q in enumerate(raw): if q is None: raise ValueError( f"CustomRecipe qfactory returned None for slot {i} " f"(role={roles[i]}). Every slot must return a Quantizer " - "instance or a QuantizerRequest." + "instance or a QuantizerRequest. For an intentional no-op " + "(high-precision / unquantized) slot, return an " + "IdentityQuantizer instead of None." ) # -- Delayed scaling sub-state -- diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index a167af4fa7..3bc1632085 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -9,7 +9,7 @@ from __future__ import annotations from torch.utils.cpp_extension import IS_HIP_EXTENSION import os -from typing import Optional, Tuple, Iterable, Any, Dict, Union +from typing import NamedTuple, Optional, Tuple, Iterable, Any, Dict, Union, get_type_hints import abc import warnings import math @@ -27,13 +27,33 @@ _stride_from_shape, ) - # Custom ops that should pass through __torch_dispatch__ without unwrapping # QuantizedTensor subclasses (e.g. Float8Tensor). Register ops here that # handle quantized tensors internally. _quantized_tensor_passthrough_ops: set = set() +class InnerTensor(NamedTuple): + """Marks a storage field as a flat inner tensor. + + Annotate the field with it -- ``_scale_inv: Annotated[torch.Tensor, + InnerTensor("fp8_scale_inv")]`` -- and ``__init_subclass__`` collects the + declarations into ``_INNER_TENSORS``, in field order. + """ + + ctor_kwarg: str + + +def _collect_inner_tensor_fields(cls: type) -> Tuple[Tuple[str, str], ...]: + """Buffers a storage class declares, as ``(attribute, constructor kwarg)``.""" + fields = [] + for attr, hint in get_type_hints(cls, include_extras=True).items(): + for meta in getattr(hint, "__metadata__", ()): + if isinstance(meta, InnerTensor): + fields.append((attr, meta.ctor_kwarg)) + return tuple(fields) + + class QuantizedTensorStorage: r"""Base class for all TensorStorage classes. @@ -53,6 +73,20 @@ class QuantizedTensorStorage: _dtype: torch.dtype _quantizer: Optional[Quantizer] + @property + def shape(self) -> torch.Size: + """Logical tensor shape, valid on bare storages and wrapper tensors alike. + + Wrapper subclasses (also ``torch.Tensor``) defer to the native tensor + shape (``size()`` on a bare storage may reconstruct the shape from the + columnwise buffer, which is not necessarily the outer shape); bare + storages derive it from ``size()``. + """ + if isinstance(self, torch.Tensor): + # pylint: disable=unnecessary-dunder-call + return torch._C.TensorBase.shape.__get__(self, type(self)) + return torch.Size(self.size()) + def update_usage( self, rowwise_usage: Optional[bool] = None, @@ -137,6 +171,120 @@ def copy_from_storage(self, src: QuantizedTensorStorage) -> None: f"{self.__class__.__name__} class does not implement copy_from_storage function" ) + # ── FSDP2 buffer protocol ─────────────────────────────────────── + # + # These three methods decouple FSDP2 all-gather buffer extraction from + # format-specific padding/layout tricks. `HybridQuantizedTensor` uses them + # to aggregate buffers from its two sub-storages without knowing each + # sub-storage's internal field layout. + # + # Contract: + # * ``fsdp_buffer_fields`` returns an ordered tuple of attribute names + # on *self* that hold the tensor buffers that must be all-gathered. + # Scalars/metadata that only need broadcasting (e.g. per-tensor FP8 + # ``_scale_inv``) are NOT listed here — they travel via the hook's + # metadata tuple instead. + # * ``fsdp_extract_buffers`` returns ``(buffers, reassembly_meta)``. + # The default implementation reads the fields as-is. Sub-storages with + # gather-time padding (MXFP8 block scales) override this to strip the + # padding before gather. + # * ``fsdp_assign_gathered`` writes the gathered buffers back into the + # storage's fields. Sub-storages with gather-time padding override + # this to re-apply the padding before assignment. + + def fsdp_buffer_fields(self) -> Tuple[str, ...]: + """Ordered attribute names holding tensor buffers gathered by FSDP2.""" + raise NotImplementedError( + f"{self.__class__.__name__} class does not implement fsdp_buffer_fields" + ) + + def fsdp_extract_buffers( + self, + ) -> Tuple[Tuple[Optional[torch.Tensor], ...], Dict[str, Any]]: + """Return ``(buffers, reassembly_meta)`` for FSDP2 all-gather. + + Default implementation reads the fields named by ``fsdp_buffer_fields`` + verbatim. Override when the on-disk layout differs from the + gather-ready layout (e.g. MXFP8 block scales carry alignment padding). + """ + names = self.fsdp_buffer_fields() + buffers = tuple(getattr(self, name) for name in names) + return buffers, {"field_names": names} + + def fsdp_assign_gathered( + self, + gathered: Tuple[Optional[torch.Tensor], ...], + meta: Dict[str, Any], + ) -> None: + """Write gathered buffers into the fields named in ``meta``. + + Override when the gather-ready layout needs a format-specific transform + (e.g. MXFP8 scales must be padded back to ``[128, 4]`` / ``[4, 128]``). + """ + names = meta["field_names"] + if len(names) != len(gathered): + raise RuntimeError( + f"{type(self).__name__}.fsdp_assign_gathered got " + f"{len(gathered)} buffers for {len(names)} fields" + ) + for name, buf in zip(names, gathered): + setattr(self, name, buf) + + # ----- PyTorch subclass flatten protocol (torch.compile / TensorSpec) ----- + + # Collected from the subclasses' :class:`InnerTensor` field annotations; everything + # else returned by :meth:`get_metadata` is treated as non-tensor context. + _INNER_TENSORS: Tuple[Tuple[str, str], ...] = () + + def __init_subclass__(cls, **kwargs) -> None: + super().__init_subclass__(**kwargs) + cls._INNER_TENSORS = _collect_inner_tensor_fields(cls) + + def _flatten_nontensor_kwargs(self) -> Dict[str, Any]: + """Non-tensor constructor kwargs (scalars, dtype, quantizer).""" + tensor_kwargs = {kwarg for _, kwarg in self._INNER_TENSORS} + return {k: v for k, v in self.get_metadata().items() if k not in tensor_kwargs} + + def __tensor_flatten__(self) -> Tuple[list, Dict[str, Any]]: + """Return ``(inner_tensor_attr_names, context)``; see class comment.""" + present = [attr for attr, _ in self._INNER_TENSORS if getattr(self, attr) is not None] + ctx = { + "cls": type(self), + "is_tensor": isinstance(self, QuantizedTensor), + "requires_grad": ( + bool(self.requires_grad) if isinstance(self, QuantizedTensor) else False + ), + "nontensor_kwargs": self._flatten_nontensor_kwargs(), + } + return present, ctx + + @staticmethod + def __tensor_unflatten__( + inner_tensors: Dict[str, torch.Tensor], + ctx: Dict[str, Any], + outer_size: Iterable[int], + outer_stride: Optional[Iterable[int]], + ) -> QuantizedTensorStorage: + """Rebuild a storage / wrapper from flat tensors + context.""" + cls = ctx["cls"] + kwargs: Dict[str, Any] = dict(ctx["nontensor_kwargs"]) + # Map each declared inner tensor back to its constructor kwarg (absent -> None). + for attr, kwarg in cls._INNER_TENSORS: + kwargs[kwarg] = inner_tensors.get(attr) + if not ctx["is_tensor"]: + return cls(**kwargs) + # Wrapper subclass: it also needs outer shape / dtype / device / stride. + fake_dtype = kwargs.get("fake_dtype") + device = next((t.device for t in inner_tensors.values() if t is not None), None) + return cls( + shape=tuple(outer_size), + dtype=fake_dtype, + requires_grad=ctx["requires_grad"], + device=device, + stride=tuple(outer_stride) if outer_stride is not None else None, + **kwargs, + ) + def prepare_for_saving( *tensors: Union[torch.Tensor, QuantizedTensorStorage], @@ -354,6 +502,72 @@ def make_empty( result.requires_grad_(True) return result + # ----- Data-free inner-tensor/metadata primitives backing TensorSpec ----- + + def inner_tensor_specs( + self, shape: Tuple[int, ...] + ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]: + """Return ``{attr_name: (shape, dtype)}`` for the inner tensors + this quantizer would allocate for a logical tensor of ``shape``. + + Keys must match the inner-tensor attribute names declared in the storage's + ``_INNER_TENSORS``, be emitted in that same order, and respect the + quantizer's usage flags. + """ + raise NotImplementedError( + f"{self.__class__.__name__} does not implement inner_tensor_specs; " + "it cannot be used with TensorSpec / pure-Python allocation" + ) + + def storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + """Non-tensor context for the produced storage. + + Returns ``{"cls": , "nontensor_kwargs": {...}}`` where ``cls`` is + the concrete class to instantiate (wrapper subclass for user-visible + tensors, bare storage class for ``internal`` quantizers) and + ``nontensor_kwargs`` are its non-tensor constructor kwargs (e.g. + ``fp8_dtype``, ``quantizer``, ``fake_dtype``). + """ + raise NotImplementedError( + f"{self.__class__.__name__} does not implement storage_metadata; " + "it cannot be used with TensorSpec / pure-Python allocation" + ) + + def alloc_tensors( + self, + shape: Iterable[int], + *, + device: Optional[Union[torch.device, str]] = None, + ) -> Dict[str, torch.Tensor]: + """Allocate (uninitialized) the flat inner tensors for ``shape``. + + Returns ``{attr_name: torch.Tensor}`` suitable as the ``inner_tensors`` + argument of the storage's ``__tensor_unflatten__``. + """ + device = torch.device(device if device is not None else "cuda") + return { + attr: torch.empty(buf_shape, dtype=buf_dtype, device=device) + for attr, (buf_shape, buf_dtype) in self.inner_tensor_specs(tuple(shape)).items() + } + + def create_metadata( + self, + _shape: Iterable[int], + *, + dtype: torch.dtype, + requires_grad: bool = False, + ) -> Dict[str, Any]: + """Build the data-free ``__tensor_unflatten__`` context describing the + quantized tensor this quantizer would produce for ``shape`` / ``dtype``. + """ + meta = self.storage_metadata(dtype) + return { + "cls": meta["cls"], + "is_tensor": not self.internal, + "requires_grad": requires_grad, + "nontensor_kwargs": meta["nontensor_kwargs"], + } + def calibrate(self, tensor: torch.Tensor) -> None: """Calibrate quantizer state @@ -398,6 +612,14 @@ def supports_only_rowwise_all_gather(self) -> bool: """Returns True if the quantizer supports only rowwise all-gather""" return False + def is_requantization_safe(self) -> bool: + """Whether repeated quantization of the same input reproduces the same value. + + Stateful or stochastic quantizers should return ``False``. This lets callers + decide whether a quantized value may be discarded and reconstructed later. + """ + return False + def is_quantizable(self, inp: torch.Tensor) -> bool: # pylint: disable=unused-argument """Whether tensor supports quantized all-gather @@ -505,6 +727,7 @@ def __new__( requires_grad: bool = False, device: Optional[torch.device] = None, stride: Optional[Iterable[int]] = None, + storage_offset: int = 0, ): if fake_dtype is not None and fake_dtype != dtype: raise ValueError(f"fake_dtype ({fake_dtype}) does not match dtype ({dtype})") @@ -523,7 +746,7 @@ def __new__( cls, shape, strides=stride, - storage_offset=0, + storage_offset=storage_offset, dtype=dtype, layout=torch.strided, requires_grad=requires_grad, @@ -718,7 +941,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): # View op if func == torch.ops.aten.view.default: - raise NotImplementedError("{cls.__name__} class does not support tensor views") + raise NotImplementedError(f"{cls.__name__} class does not support tensor views") # New empty op (used by DCP async staging to create CPU copies) if func == torch.ops.aten.new_empty.default: diff --git a/transformer_engine/pytorch/tensor/__init__.py b/transformer_engine/pytorch/tensor/__init__.py index 426c656d47..c3355b6c62 100644 --- a/transformer_engine/pytorch/tensor/__init__.py +++ b/transformer_engine/pytorch/tensor/__init__.py @@ -19,11 +19,15 @@ from .storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from .storage.nvfp4_tensor_storage import NVFP4TensorStorage from .storage.grouped_tensor_storage import GroupedTensorStorage +from .storage.hybrid_tensor_storage import HybridQuantizedTensorStorage from .float8_tensor import Float8Tensor, Float8Quantizer, Float8CurrentScalingQuantizer from .mxfp8_tensor import MXFP8Tensor, MXFP8Quantizer from .float8_blockwise_tensor import Float8BlockwiseQTensor, Float8BlockQuantizer from .nvfp4_tensor import NVFP4Tensor, NVFP4Quantizer from .grouped_tensor import GroupedTensor +from .hybrid_tensor import HybridQuantizedTensor, HybridQuantizer +from .identity_tensor import IdentityTensor, IdentityQuantizer +from .storage.identity_tensor_storage import IdentityTensorStorage from .utils import cast_master_weights_to_fp8, replace_raw_data __all__ = [ @@ -33,18 +37,24 @@ "MXFP8Quantizer", "Float8BlockQuantizer", "NVFP4Quantizer", + "HybridQuantizer", + "IdentityQuantizer", "QuantizedTensorStorage", "Float8TensorStorage", "MXFP8TensorStorage", "Float8BlockwiseQTensorStorage", "NVFP4TensorStorage", "GroupedTensorStorage", + "HybridQuantizedTensorStorage", + "IdentityTensorStorage", "QuantizedTensor", "Float8Tensor", + "IdentityTensor", "MXFP8Tensor", "Float8BlockwiseQTensor", "NVFP4Tensor", "GroupedTensor", + "HybridQuantizedTensor", "prepare_for_saving", "restore_from_saved", "restore_from_func_ctx", @@ -97,5 +107,9 @@ def get_all_tensor_types(): NVFP4TensorStorage, GroupedTensor, GroupedTensorStorage, + HybridQuantizedTensor, + HybridQuantizedTensorStorage, + IdentityTensor, + IdentityTensorStorage, ] return all_tensor_types diff --git a/transformer_engine/pytorch/tensor/_quantization_helpers.py b/transformer_engine/pytorch/tensor/_quantization_helpers.py index 1b08039dda..10672bbcfb 100644 --- a/transformer_engine/pytorch/tensor/_quantization_helpers.py +++ b/transformer_engine/pytorch/tensor/_quantization_helpers.py @@ -9,13 +9,66 @@ """ from __future__ import annotations -from typing import Callable, Optional, Tuple, Any, Dict, TYPE_CHECKING +from typing import Callable, Optional, Tuple, Any, Dict, Iterable, TYPE_CHECKING import torch if TYPE_CHECKING: from transformer_engine.pytorch.quantized_tensor import QuantizedTensor +def _resolve_view_shape(input_shape: Iterable[int], shape: Iterable[int]) -> torch.Size: + """Resolve a requested view shape with PyTorch-compatible semantics. + + The concrete-integer path avoids constructing a temporary meta tensor. If + either shape contains symbolic dimensions, retain the previous meta-tensor + path so that PyTorch remains responsible for symbolic shape handling. + """ + input_shape = tuple(input_shape) + shape = tuple(shape) + if len(shape) == 1 and isinstance(shape[0], (list, tuple, torch.Size)): + shape = tuple(shape[0]) + + # Avoid comparisons that specialize or guard SymInts. The meta fallback is + # also useful for preserving PyTorch's type checking of non-integer dims. + if any(not isinstance(dim, int) or isinstance(dim, bool) for dim in (*input_shape, *shape)): + return torch.empty(input_shape, device="meta").view(shape).shape + + input_numel = 1 + for dim in input_shape: + input_numel *= dim + + inferred_dim = None + known_numel = 1 + for index, dim in enumerate(shape): + if dim == -1: + if inferred_dim is not None: + raise RuntimeError("only one dimension can be inferred") + inferred_dim = index + elif dim < 0: + raise RuntimeError( + f"invalid shape dimension {dim} at index {index} of shape {list(shape)}" + ) + else: + known_numel *= dim + + resolved_shape = list(shape) + if inferred_dim is not None: + if known_numel == 0: + if input_numel == 0: + raise RuntimeError( + f"cannot reshape tensor of 0 elements into shape {list(shape)} because " + "the unspecified dimension size -1 can be any value and is ambiguous" + ) + raise RuntimeError(f"shape '{list(shape)}' is invalid for input of size {input_numel}") + if input_numel % known_numel != 0: + raise RuntimeError(f"shape '{list(shape)}' is invalid for input of size {input_numel}") + resolved_shape[inferred_dim] = input_numel // known_numel + elif known_numel != input_numel: + raise RuntimeError(f"shape '{list(shape)}' is invalid for input of size {input_numel}") + + return torch.Size(resolved_shape) + + class _QuantizeFunc(torch.autograd.Function): """Quantize tensor""" diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index 4126db32a6..a56f5d2b59 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -9,7 +9,7 @@ from collections.abc import Iterable import math import warnings -from typing import Any, Optional, Tuple, Union +from typing import Any, Dict, Optional, Tuple, Union import torch import transformer_engine_torch as tex @@ -77,6 +77,39 @@ def copy(self) -> Float8BlockQuantizer: return quantizer + # ----- TensorSpec / pure-Python allocation ----- + + def storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + return { + "cls": Float8BlockwiseQTensorStorage if self.internal else Float8BlockwiseQTensor, + "nontensor_kwargs": { + "fp8_dtype": self.dtype, + "quantizer": self, + "is_2D_scaled": self.block_scaling_dim == 2, + "fake_dtype": fake_dtype, + }, + } + + def inner_tensor_specs( + self, shape: Tuple[int, ...] + ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]: + shape = tuple(shape) + specs: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {} + # Blockwise FP8 scales are FP32; columnwise data is stored transposed. + if self.rowwise_usage: + specs["_rowwise_data"] = (shape, torch.uint8) + specs["_rowwise_scale_inv"] = ( + tuple(self.get_scale_shape(shape, columnwise=False)), + torch.float32, + ) + if self.columnwise_usage: + specs["_columnwise_data"] = (tuple(self.get_columnwise_shape(shape)), torch.uint8) + specs["_columnwise_scale_inv"] = ( + tuple(self.get_scale_shape(shape, columnwise=True)), + torch.float32, + ) + return specs + def update_quantized( self, src: torch.Tensor, @@ -126,6 +159,10 @@ def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: """Quantize tensor implementation""" return tex.quantize(tensor, self) + def is_requantization_safe(self) -> bool: + """Block-FP8 scales are derived deterministically from each input.""" + return True + def get_scale_shape(self, shape: Iterable[int], columnwise: bool) -> Tuple[int, int]: """Scaling tensor shape. @@ -563,7 +600,11 @@ def shape(self): if self._rowwise_data is not None: return self._rowwise_data.shape if self._columnwise_data is not None: - return self._columnwise_data.shape + # Columnwise data is stored transposed, matching size() in the storage. + dims = self._columnwise_data.shape + if len(dims) == 2: + return torch.Size((dims[1], dims[0])) + return torch.Size(tuple(dims[1:]) + (dims[0],)) return torch.Tensor.size(self) @property diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 5c2d42441c..8dbb17618f 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -6,8 +6,7 @@ """Tensor class with FP8 data""" from __future__ import annotations - -from typing import Any, Optional, Tuple, Iterable, Union +from typing import Any, Dict, Optional, Tuple, Iterable, Union import warnings import torch from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState @@ -18,11 +17,15 @@ Float8CurrentScaling, Recipe, ) -from ..utils import canonicalize_process_group, devices_match +from ..utils import canonicalize_process_group, devices_match, is_non_tn_fp8_gemm_supported from .storage.float8_tensor_storage import Float8TensorStorage, _FromFloat8Func from ..quantized_tensor import QuantizedTensor, Quantizer from ..dynamo import register_value_opaque_quantizer -from ._quantization_helpers import _IdentityFunc, safe_quantized_repr +from ._quantization_helpers import ( + _IdentityFunc, + _resolve_view_shape, + safe_quantized_repr, +) from ..constants import dist_group_type, DType from torch.utils.cpp_extension import IS_HIP_EXTENSION @@ -46,6 +49,14 @@ } +def _columnwise_shape_for(rowwise_shape: Iterable[int]) -> torch.Size: + """Physical columnwise FP8 shape for a logical rowwise shape.""" + shape = torch.Size(rowwise_shape) + if len(shape) == 0: + return shape + return torch.Size((shape[-1], *shape[:-1])) + + class Float8Quantizer(Quantizer): """Builder class for FP8 tensors with per-tensor delayed scaling @@ -417,6 +428,40 @@ def supports_only_rowwise_all_gather(self) -> bool: """ return True + # ----- TensorSpec / pure-Python allocation ----- + + def storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + return { + "cls": Float8TensorStorage if self.internal else Float8Tensor, + "nontensor_kwargs": { + "fp8_dtype": self.dtype, + "quantizer": self, + "fake_dtype": fake_dtype, + }, + } + + def inner_tensor_specs( + self, shape: Tuple[int, ...] + ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]: + shape = tuple(shape) + specs: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {} + # Mirror the C++ quantizer allocation (csrc/quantizer.cpp): on non-TN-capable + # archs (Blackwell+) a single ``_data`` buffer backs both row- and column-wise + # usage and no separate transpose is materialized. This must match what the + # real kernel produces so the torch.compile fake layout lines up slot-for-slot. + non_tn = is_non_tn_fp8_gemm_supported() + if self.rowwise_usage or non_tn: + specs["_data"] = (shape, torch.uint8) + if self.columnwise_usage and not non_tn: + specs["_transpose"] = ((shape[-1], *shape[:-1]), torch.uint8) + # Per-tensor scale-inv is always present for current scaling. + specs["_scale_inv"] = ((1,), torch.float32) + return specs + + def is_requantization_safe(self) -> bool: + """Current scaling is derived deterministically from each input.""" + return True + register_value_opaque_quantizer(Float8CurrentScalingQuantizer) @@ -509,8 +554,12 @@ def detach(self) -> Float8Tensor: def clone(self) -> Float8Tensor: # pylint: disable=missing-function-docstring - assert self._data is not None - data = self._data.detach().clone() + # ``_data`` may be None for columnwise-only sub-storages of a + # HybridQuantizedTensor on architectures without native non-TN FP8 + # GEMM (Hopper / L40), where columnwise-only Float8 allocates + # ``_transpose`` instead of ``_data``. On Blackwell+ the C++ + # override keeps ``_data`` populated even in columnwise-only mode. + data = self._data.detach().clone() if self._data is not None else None data_transpose = None if self._transpose is not None: data_transpose = self._transpose.detach().clone() @@ -567,6 +616,12 @@ def _reset_caches(self) -> None: Set transpose cache as invalid. Should be called after any in-place operation. """ + if self._data is None and self._transpose is not None: + # Columnwise-only Float8 tensors on Hopper / L40 store their only + # live FP8 payload in _transpose. Treat it as primary storage, not + # as a derived cache that can be invalidated. + self._transpose_invalid = False + return self._transpose_invalid = True def remove_caches(self) -> None: @@ -611,24 +666,34 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): if func == aten.view.default: tensor = args[0] data = tensor._data - out_data = data.__torch_dispatch__( - func, - types, - [data] + list(args[1:]), - kwargs, - ) - out_shape = out_data.size() + out_data = None + if data is not None: + out_data = data.__torch_dispatch__( + func, + types, + [data] + list(args[1:]), + kwargs, + ) + out_shape = out_data.size() + else: + out_shape = _resolve_view_shape(tensor.shape, args[1:]) + out_transpose = None if tensor._transpose_invalid else tensor._transpose if out_transpose is not None: - out_transpose_shape = out_transpose.size() - if ( - out_transpose_shape[0] != out_shape[-1] - or out_transpose_shape[1:] != out_shape[:-1] - ): + view_shape_for_transpose = _columnwise_shape_for(out_shape) + if out_transpose.shape != view_shape_for_transpose: + if data is None: + raise NotImplementedError( + "Float8Tensor view with columnwise-only data is only supported " + "when the requested shape preserves the columnwise layout" + ) out_transpose = None else: - view_shape_for_transpose = [out_shape[-1]] + list(out_shape[:-1]) out_transpose = out_transpose.view(*view_shape_for_transpose) + if data is None and out_transpose is None: + raise NotImplementedError( + "Float8Tensor view with columnwise-only data requires a valid columnwise buffer" + ) return Float8Tensor( shape=out_shape, dtype=tensor.dtype, @@ -644,17 +709,22 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): if func in (aten.slice.Tensor, aten.select.int): tensor = args[0] data = tensor._data - data_slice = data.__torch_dispatch__( - func, - types, - [data] + list(args[1:]), - kwargs, - ) + data_slice = None + if data is not None: + data_slice = data.__torch_dispatch__( + func, + types, + [data] + list(args[1:]), + kwargs, + ) transpose_slice = None if tensor._transpose is not None and not tensor._transpose_invalid: transpose = tensor._transpose - ndim = data.dim() + ndim = tensor.dim() + if ndim == 0: + return super().__torch_dispatch__(func, types, args, kwargs) dim = args[1] if len(args) > 1 else 0 + dim %= ndim t_dim = 0 if dim == ndim - 1 else dim + 1 transpose_slice = transpose.__torch_dispatch__( func, @@ -662,34 +732,55 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): [transpose, t_dim] + list(args[2:]), kwargs, ) + if func == aten.select.int and dim == ndim - 1 and transpose_slice.dim() > 1: + transpose_slice = transpose_slice.movedim(-1, 0).contiguous() + + if data_slice is not None: + out_shape = data_slice.shape + else: + logical = torch.empty(tensor.shape, device="meta") + out_shape = func(logical, *args[1:], **(kwargs or {})).shape + if transpose_slice is None: + raise RuntimeError( + "Float8Tensor slice/select requires rowwise or columnwise data" + ) + expected_transpose_shape = _columnwise_shape_for(out_shape) + if transpose_slice.shape != expected_transpose_shape: + raise RuntimeError( + "Float8Tensor slice/select produced incompatible columnwise storage: " + f"expected {tuple(expected_transpose_shape)}, got " + f"{tuple(transpose_slice.shape)}" + ) return Float8Tensor.make_like( tensor, data=data_slice, data_transpose=transpose_slice, - shape=data_slice.shape, + shape=out_shape, ) # Related to FSDP2 if func == aten.split.Tensor: tensor = args[0] data = tensor._data - func_out = data.__torch_dispatch__( - func, - types, - [data] + list(args[1:]), - kwargs, - ) - t_func_out = [None] * len(func_out) - # Compute corresponding split of the transpose cache if available + # _data may be None for columnwise-only sub-storages (hybrid quantization) + if data is not None: + func_out = data.__torch_dispatch__( + func, + types, + [data] + list(args[1:]), + kwargs, + ) + else: + func_out = None + + t_func_out = None if tensor._transpose is not None and not tensor._transpose_invalid: transpose = tensor._transpose - ndim = data.dim() - # Figure out the original split dim + ndim = tensor.dim() if "dim" in kwargs: dim_to_split = kwargs["dim"] else: dim_to_split = args[2] if len(args) > 2 else 0 - # Dimension along which transpose needs to be split t_dim = 0 if dim_to_split == ndim - 1 else dim_to_split + 1 t_func_out = transpose.__torch_dispatch__( func, @@ -697,12 +788,30 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): [transpose, args[1], t_dim], kwargs, ) + + ref_out = func_out if func_out is not None else t_func_out + if ref_out is None: + return super().__torch_dispatch__(func, types, args, kwargs) + + num_splits = len(ref_out) + if func_out is None: + func_out = [None] * num_splits + if t_func_out is None: + t_func_out = [None] * num_splits + outs = [ Float8Tensor.make_like( tensor, data=split_tensor, data_transpose=split_transpose_tensor, - shape=split_tensor.shape, + shape=( + split_tensor.shape + if split_tensor is not None + else ( + *split_transpose_tensor.shape[1:], + split_transpose_tensor.shape[0], + ) + ), ) for split_tensor, split_transpose_tensor in zip(func_out, t_func_out) ] @@ -712,12 +821,18 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): # create fresh new tensor with zeros. tensor = args[0] data = tensor._data - func_out = data.__torch_dispatch__( - func, - types, - [data] + list(args[1:]), - kwargs, - ) + storage_kwargs = dict(kwargs or {}) + output_dtype = storage_kwargs.pop("dtype", None) or tensor.dtype + storage_kwargs.pop("layout", None) + storage_kwargs.pop("requires_grad", None) + func_out = None + if data is not None: + func_out = data.__torch_dispatch__( + func, + types, + [data] + list(args[1:]), + storage_kwargs, + ) func_transposed_out = None if tensor._transpose is not None and not tensor._transpose_invalid: transpose = tensor._transpose @@ -727,25 +842,69 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): func, types, [transpose, t_shape] + list(args[2:]), - kwargs, + storage_kwargs, ) + if func_out is None and func_transposed_out is None: + raise RuntimeError("Float8Tensor.new_zeros requires rowwise or columnwise data") scale_inv = tensor._scale_inv.detach().clone() + reference = func_out if func_out is not None else func_transposed_out + if scale_inv.device != reference.device: + scale_inv = scale_inv.to(reference.device) quantizer = tensor._quantizer # Deep-copied in constructor out_tensor = Float8Tensor( data=func_out, - shape=func_out.shape, - dtype=tensor.dtype, + shape=torch.Size(args[1]), + dtype=output_dtype, fp8_dtype=tensor._fp8_dtype, fp8_scale_inv=scale_inv, data_transpose=func_transposed_out, quantizer=quantizer, - device=tensor.device, + device=reference.device, ) return out_tensor if func == torch.ops.aten.as_strided.default: tensor = args[0] data = tensor._data + if data is None: + size = torch.Size(args[1]) + stride = tuple(args[2]) + storage_offset = (kwargs or {}).get( + "storage_offset", + args[3] if len(args) > 3 else tensor.storage_offset(), + ) + # A contiguous 2D row shard maps to a column slice in the + # persistent transpose and can remain a Float8Tensor. + has_valid_transpose = ( + tensor._transpose is not None and not tensor._transpose_invalid + ) + is_contiguous_row_shard = ( + tensor.dim() == 2 + and tensor.shape[1] > 0 + and len(size) == 2 + and size[1] == tensor.shape[1] + and stride == tuple(tensor.stride()) + ) + if ( + has_valid_transpose + and is_contiguous_row_shard + and storage_offset % tensor.shape[1] == 0 + ): + row_start = storage_offset // tensor.shape[1] + row_end = row_start + size[0] + if 0 <= row_start and row_end <= tensor.shape[0]: + transpose_out = tensor._transpose[:, row_start:row_end] + return Float8Tensor.make_like( + tensor, + data=None, + data_transpose=transpose_out, + shape=size, + ) + + # Arbitrary logical strides are not generally affine views of + # transposed storage. Fall back explicitly to high precision. + return func(tensor.dequantize(), *args[1:], **(kwargs or {})) + # Apply as_strided to the primary uint8 data func_out = data.__torch_dispatch__( func, @@ -957,9 +1116,19 @@ def __reduce_ex__(self, protocol: int) -> tuple: with ``torch.serialization.add_safe_globals`` to keep ``torch.load(weights_only=True)`` compatibility. """ + data_transpose = None + if self._data is None and self._transpose is not None and not self._transpose_invalid: + data_transpose = self._transpose return ( _make_float8_tensor_in_reduce_ex, - (self._data, self._fp8_dtype, self._scale_inv, self.dtype, self.shape), + ( + self._data, + self._fp8_dtype, + self._scale_inv, + self.dtype, + self.shape, + data_transpose, + ), ) @classmethod @@ -1055,11 +1224,12 @@ def _set_data(self, tensor: torch.Tensor) -> None: def _make_float8_tensor_in_reduce_ex( - data: torch.Tensor, + data: Optional[torch.Tensor], fp8_dtype: DType, fp8_scale_inv: torch.Tensor, dtype: torch.dtype, shape: torch.Size, + data_transpose: Optional[torch.Tensor] = None, ) -> Float8Tensor: """Reconstruct a ``Float8Tensor`` from its ``__reduce_ex__`` payload.""" return Float8Tensor( @@ -1068,7 +1238,12 @@ def _make_float8_tensor_in_reduce_ex( fp8_scale_inv=fp8_scale_inv, dtype=dtype, shape=shape, - device=data.device if data is not None else None, + data_transpose=data_transpose, + device=( + data.device + if data is not None + else data_transpose.device if data_transpose is not None else None + ), ) @@ -1089,16 +1264,28 @@ def forward( ctx.shape = tensor.shape if shape is None: return tensor.detach() - out_data = tensor._data.view(*shape) - out_shape = out_data.size() + out_data = None + if tensor._data is not None: + out_data = tensor._data.view(*shape) + out_shape = out_data.size() + else: + out_shape = _resolve_view_shape(tensor.shape, shape) out_transpose = None if tensor._transpose_invalid else tensor._transpose if out_transpose is not None: - out_transpose_shape = out_transpose.size() - if out_transpose_shape[0] != out_shape[-1] or out_transpose_shape[1:] != out_shape[:-1]: + view_shape_for_transpose = _columnwise_shape_for(out_shape) + if out_transpose.shape != view_shape_for_transpose: + if tensor._data is None: + raise NotImplementedError( + "Float8Tensor view with columnwise-only data is only supported " + "when the requested shape preserves the columnwise layout" + ) out_transpose = None else: - view_shape_for_transpose = [shape[-1]] + list(shape[:-1]) out_transpose = out_transpose.view(*view_shape_for_transpose) + if tensor._data is None and out_transpose is None: + raise NotImplementedError( + "Float8Tensor view with columnwise-only data requires a valid columnwise buffer" + ) return Float8Tensor( shape=out_shape, dtype=tensor.dtype, diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py new file mode 100644 index 0000000000..dc65c9894b --- /dev/null +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -0,0 +1,1169 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tensor and quantizer classes for composed rowwise and columnwise representations.""" + +from __future__ import annotations +from typing import Any, Dict, Iterable, Literal, Optional, Tuple, Union + +import torch + +from .storage.hybrid_tensor_storage import HybridQuantizedTensorStorage +from .storage.identity_tensor_storage import IdentityTensorStorage +from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer + +aten = torch.ops.aten + + +class HybridQuantizer(Quantizer): + """Quantizer that composes rowwise and columnwise representations. + + When both representations are requested, applies ``rowwise_quantizer`` to + produce the rowwise representation and ``columnwise_quantizer`` to produce + the columnwise representation. The results are wrapped in a + ``HybridQuantizedTensor``. The children may use the same or different + Transformer Engine formats, or custom ``Quantizer`` implementations. + + Parameters + ---------- + rowwise_quantizer : Quantizer + Quantizer for the rowwise direction (e.g. MXFP8Quantizer). + columnwise_quantizer : Quantizer + Quantizer for the columnwise direction (e.g. NVFP4Quantizer). + columnwise_source : {"original", "rowwise_dequantized"}, default = "original" + Source tensor for columnwise quantization. ``"original"`` quantizes + columnwise directly from the input tensor. ``"rowwise_dequantized"`` + quantizes rowwise first, dequantizes the rowwise result, then uses that + value as the columnwise source. + + Notes + ----- + Rowwise and columnwise describe storage and GEMM orientations. Whether a + representation is consumed in the forward or backward pass depends on the + tensor's role in the operation. + + ``HybridQuantizer`` pins each sub-quantizer to its designated direction by + mutating its usage flags, so it takes ownership of the supplied quantizer + instances. The rowwise and columnwise quantizers must be distinct objects. + If both directions need shared state, construct two quantizer instances that + reference the same external state object. + + Each ``HybridQuantizer`` must receive its own sub-quantizer instances; do not + reuse a sub-quantizer instance across multiple ``HybridQuantizer`` objects. + + Example + ------- + MXFP8 rowwise data plus a high-precision columnwise representation derived + from the rowwise value can be expressed as:: + + HybridQuantizer( + rowwise_quantizer=mxfp8_quantizer, + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + + In a ``CustomRecipe`` factory this can be combined with role-based + specialization. For example, a factory can return a ``HybridQuantizer`` only + for ``role.tensor_type == "weight"`` and use regular quantizers for inputs + and gradients. See + ``custom_recipes.quantizer_factory_zoo.nvfp4_1d_weight_factory`` + for a weight-only double-quantization example. + + """ + + _COLUMNWISE_SOURCES = ("original", "rowwise_dequantized") + + rowwise_quantizer: Quantizer + columnwise_quantizer: Quantizer + columnwise_source: Literal["original", "rowwise_dequantized"] + + def __init__( + self, + *, + rowwise_quantizer: Quantizer, + columnwise_quantizer: Quantizer, + columnwise_source: Literal["original", "rowwise_dequantized"] = "original", + ) -> None: + super().__init__(rowwise=True, columnwise=True) + from transformer_engine.pytorch.quantization import QuantizerRequest # local import + + for role, quantizer in ( + ("rowwise", rowwise_quantizer), + ("columnwise", columnwise_quantizer), + ): + if isinstance(quantizer, QuantizerRequest): + # TODO(#3158): Support delayed-scaling requests inside hybrid sub-quantizers. + raise TypeError( + "HybridQuantizer does not support nested QuantizerRequest " + f"objects yet; got {type(quantizer).__name__} for the {role} " + "direction. Delayed scaling in CustomRecipe is currently " + "supported only when the qfactory returns DelayedScalingRequest " + "as a top-level slot. Resolving delayed-scaling requests inside " + "HybridQuantizer is future work; pass a concrete Quantizer " + "instance instead." + ) + if not isinstance(quantizer, Quantizer): + raise TypeError( + "HybridQuantizer requires concrete Quantizer instances for " + f"both directions, but the {role} argument is " + f"{type(quantizer).__name__}." + ) + if rowwise_quantizer is columnwise_quantizer: + raise ValueError( + "HybridQuantizer requires distinct rowwise and columnwise quantizer" + " instances. If both directions need shared state, construct two" + " quantizer objects that reference the same shared state." + ) + if columnwise_source not in self._COLUMNWISE_SOURCES: + raise ValueError( + "HybridQuantizer columnwise_source must be one of " + f"{self._COLUMNWISE_SOURCES}, got {columnwise_source!r}." + ) + self.rowwise_quantizer = rowwise_quantizer + self.columnwise_quantizer = columnwise_quantizer + self.columnwise_source = columnwise_source + + # Pin each sub-quantizer to its designated direction + self.rowwise_quantizer.set_usage(rowwise=True, columnwise=False) + self.columnwise_quantizer.set_usage(rowwise=False, columnwise=True) + + def __repr__(self): + return ( + f"{self.__class__.__name__}(" + f"rowwise_usage={self.rowwise_usage}, " + f"columnwise_usage={self.columnwise_usage}, " + f"columnwise_source={self.columnwise_source!r}, " + f"internal={self.internal}, " + ")" + ) + + def copy(self) -> "HybridQuantizer": + """Create a shallow copy, preserving parent and sub-quantizer state.""" + quantizer = HybridQuantizer( + rowwise_quantizer=self.rowwise_quantizer.copy(), + columnwise_quantizer=self.columnwise_quantizer.copy(), + columnwise_source=self.columnwise_source, + ) + quantizer.set_usage( + rowwise=self.rowwise_usage, + columnwise=self.columnwise_usage, + ) + quantizer.internal = self.internal + quantizer.optimize_for_gemm = self.optimize_for_gemm + return quantizer + + @property + def with_amax_reduction(self) -> bool: + """Whether either sub-quantizer has cross-rank amax reduction enabled.""" + return getattr(self.rowwise_quantizer, "with_amax_reduction", False) or getattr( + self.columnwise_quantizer, "with_amax_reduction", False + ) + + @with_amax_reduction.setter + def with_amax_reduction(self, value: bool) -> None: + # Set on the HybridQuantizer by module / FSDP2 code, but read by the C++ + # kernel off the sub-quantizer that runs -- hence forwarded, not stored here. + for sub in (self.rowwise_quantizer, self.columnwise_quantizer): + if hasattr(sub, "with_amax_reduction"): + sub.with_amax_reduction = value + + @property + def amax_reduction_group(self): + """Amax-reduction group of the sub-quantizers, or ``None`` if unset.""" + result = None + for sub in (self.rowwise_quantizer, self.columnwise_quantizer): + group = getattr(sub, "amax_reduction_group", None) + if group is None: + continue + if result is None: + result = group + elif group is not result: + raise RuntimeError( + "HybridQuantizer sub-quantizers have inconsistent amax_reduction_group values." + ) + return result + + @amax_reduction_group.setter + def amax_reduction_group(self, value) -> None: + for sub in (self.rowwise_quantizer, self.columnwise_quantizer): + if hasattr(sub, "amax_reduction_group"): + sub.amax_reduction_group = value + + def _columnwise_src_from_rowwise( + self, + tensor: torch.Tensor, + rowwise_result: Optional[Any], + ) -> torch.Tensor: + if rowwise_result is None: + rowwise_result = self.rowwise_quantizer.quantize(tensor) + return rowwise_result.dequantize(dtype=tensor.dtype) + + def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: + # Gate each sub-quantizer call on the parent usage flag. Sub-quantizers + # are pinned to one direction in ``__init__``; the parent flag decides + # whether to invoke them. + rowwise_result = self.rowwise_quantizer.quantize(tensor) if self.rowwise_usage else None + columnwise_src = tensor + if self.columnwise_usage and self.columnwise_source == "rowwise_dequantized": + columnwise_src = self._columnwise_src_from_rowwise(tensor, rowwise_result) + columnwise_result = ( + self.columnwise_quantizer.quantize(columnwise_src) if self.columnwise_usage else None + ) + + if self.internal: + return HybridQuantizedTensorStorage( + rowwise_storage=rowwise_result, + columnwise_storage=columnwise_result, + quantizer=self, + fake_dtype=tensor.dtype, + ) + + return HybridQuantizedTensor( + shape=tensor.shape, + dtype=tensor.dtype, + rowwise_storage=rowwise_result, + columnwise_storage=columnwise_result, + quantizer=self, + ) + + def make_empty( + self, + shape: Iterable[int], + *, + dtype: torch.dtype = torch.float32, + device: Optional[torch.device] = None, + requires_grad: bool = False, + pin_memory: bool = False, + ) -> Union["HybridQuantizedTensor", HybridQuantizedTensorStorage]: + # Mirror ``quantize_impl``: invoke each sub-quantizer with its own + # ``internal`` setting (no toggle), so the produced sub-storages have + # the same type that ``quantize_impl`` would produce via + # ``sub_quantizer.quantize(tensor)``. + rowwise_empty = ( + self.rowwise_quantizer.make_empty( + shape, dtype=dtype, device=device, pin_memory=pin_memory + ) + if self.rowwise_usage + else None + ) + columnwise_empty = ( + self.columnwise_quantizer.make_empty( + shape, dtype=dtype, device=device, pin_memory=pin_memory + ) + if self.columnwise_usage + else None + ) + + if self.internal: + return HybridQuantizedTensorStorage( + rowwise_storage=rowwise_empty, + columnwise_storage=columnwise_empty, + quantizer=self, + fake_dtype=dtype, + ) + + return HybridQuantizedTensor( + shape=shape, + dtype=dtype, + requires_grad=requires_grad, + device=device, + rowwise_storage=rowwise_empty, + columnwise_storage=columnwise_empty, + quantizer=self, + ) + + def update_quantized( + self, + src: torch.Tensor, + dst: QuantizedTensorStorage, + *, + noop_flag: Optional[torch.Tensor] = None, + ) -> QuantizedTensorStorage: + """Re-quantize sub-storages of a hybrid tensor in-place. + + Each direction is refreshed only when the parent usage flag is set + **and** the corresponding sub-storage exists. + """ + if not isinstance(dst, HybridQuantizedTensorStorage): + raise ValueError( + "HybridQuantizer can only update HybridQuantizedTensorStorage, got" + f" {type(dst).__name__}" + ) + rowwise_result_for_columnwise = None + if self.rowwise_usage and dst._rowwise_storage is not None: + self.rowwise_quantizer.update_quantized(src, dst._rowwise_storage, noop_flag=noop_flag) + rowwise_result_for_columnwise = dst._rowwise_storage + if self.columnwise_usage and dst._columnwise_storage is not None: + columnwise_src = src + if self.columnwise_source == "rowwise_dequantized": + columnwise_src = self._columnwise_src_from_rowwise( + src, rowwise_result_for_columnwise + ) + self.columnwise_quantizer.update_quantized( + columnwise_src, dst._columnwise_storage, noop_flag=noop_flag + ) + return dst + + def supports_only_rowwise_all_gather(self) -> bool: + """Whether all-gather requires a rowwise-dequantizable source. + + Hybrid tensors currently use the high-precision fallback, which + dequantizes the local shard before communication. Preserve rowwise + data when required by the rowwise sub-quantizer or when the + columnwise sub-quantizer is NVFP4, whose columnwise-only storage + cannot be dequantized. + """ + if self.rowwise_quantizer.supports_only_rowwise_all_gather(): + return True + # Local import avoids a circular dependency chain + # (nvfp4_tensor → quantized_tensor → hybrid_tensor at module import). + from .nvfp4_tensor import NVFP4Quantizer # noqa: PLC0415 + + return isinstance(self.columnwise_quantizer, NVFP4Quantizer) + + def is_requantization_safe(self) -> bool: + """Whether repeated quantization reproduces all requested representations.""" + if self.rowwise_usage and not self.rowwise_quantizer.is_requantization_safe(): + return False + if self.columnwise_usage and not self.columnwise_quantizer.is_requantization_safe(): + return False + if ( + self.columnwise_usage + and self.columnwise_source == "rowwise_dequantized" + and not self.rowwise_quantizer.is_requantization_safe() + ): + return False + return True + + def _get_compatible_recipe(self): + # HybridQuantizer is only reachable via CustomRecipe (the qfactory + # returns HybridQuantizer per role). Checking that the autocast recipe + # is also CustomRecipe catches the obvious mismatch (e.g. hybrid + # quantized_model_init + built-in MXFP8BlockScaling autocast). + # We trust that users who write a CustomRecipe know what they're doing + # with regard to per-operand scaling mode compatibility. + # + # TODO(#3158): validate per-operand scaling-mode compatibility at + # recipe-build time instead of at cuBLAS-dispatch time. Concretely: + # 1. Walk the qfactory outputs for a given module_type (``linear``, + # ``grouped_linear``, ``dpa``) — call the factory for each + # ``QuantizerRole.tensor_type`` the module uses. + # 2. Extract the scaling_mode of each sub-quantizer: + # weight_row, weight_col (from HybridQuantizer) + # input_row, input_col (from HybridQuantizer) + # grad_output_row, grad_output_col (plain quantizer OR + # HybridQuantizer) + # 3. Assert the three GEMM pairs share a scaling_mode each: + # fprop TN: weight_row == input_row (FormatA) + # dgrad NN: weight_col == grad_output_row (FormatB) + # wgrad NT: input_col == grad_output_col (FormatC) + # Mismatches raise ``ValueError`` naming the offending slots, e.g. + # "dgrad GEMM: weight columnwise format (MXFP8) does not match + # grad_output rowwise format (NVFP4)". + # 4. Blocked on `semantic_quantizer_roles` / PR #2620 for the + # ``QuantizerRole`` dataclass — the factory signature is role- + # aware only on that branch. + from transformer_engine.common.recipe import CustomRecipe # avoid circular import + + return CustomRecipe + + +class HybridQuantizedTensor(HybridQuantizedTensorStorage, QuantizedTensor): + """Tensor holding independently produced rowwise and columnwise representations. + + The tensor presents as having a standard logical dtype, but + internally stores representations produced by the two sub-quantizers. + These may use the same format, different formats, or a high-precision + representation such as ``IdentityTensorStorage``. + + Parameters + ---------- + shape : iterable of int + Tensor dimensions. + dtype : torch.dtype + Nominal tensor datatype. + rowwise_storage : QuantizedTensorStorage, optional + Sub-storage for the rowwise representation. + columnwise_storage : QuantizedTensorStorage, optional + Sub-storage for the columnwise representation. + quantizer : HybridQuantizer + Parent hybrid quantizer that owns the rowwise and columnwise sub-quantizers. + requires_grad : bool, default = False + Whether to compute gradients for this tensor. + + """ + + def __new__( + cls, + *args, + rowwise_storage: Optional[QuantizedTensorStorage], + columnwise_storage: Optional[QuantizedTensorStorage], + quantizer: HybridQuantizer, + **kwargs, + ): + instance = super().__new__( + cls, + *args, + rowwise_storage=rowwise_storage, + columnwise_storage=columnwise_storage, + quantizer=quantizer, + **kwargs, + ) + return instance + + def __repr__(self, *, tensor_contents=None): + row_type = ( + type(self._rowwise_storage).__name__ if self._rowwise_storage is not None else "None" + ) + col_type = ( + type(self._columnwise_storage).__name__ + if self._columnwise_storage is not None + else "None" + ) + return ( + f"HybridQuantizedTensor(rowwise={row_type}, columnwise={col_type}, dtype={self.dtype})" + ) + + def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: + if dtype is None: + dtype = self.dtype + return HybridQuantizedTensorStorage.dequantize(self, dtype=dtype) + + def detach(self) -> HybridQuantizedTensor: + """Return a new HybridQuantizedTensor with cloned sub-storage wrappers. + + Each sub-storage is re-wrapped via its own ``make_like`` so the + new hybrid tensor has independent sub-storage objects that share + the *underlying* buffer tensors with ``self``. This is required for + the cpu_offload_v2 pattern at ``cpu_offload.py:378-382``:: + + tensor_copy = tensor.detach() + saved_tensors, _ = tensor_copy.prepare_for_saving() # nulls fields + + If ``detach()`` merely shared sub-storage objects, the + ``prepare_for_saving`` call above would null out fields on the + original ``tensor`` too (since both hybrids would point at the same + sub-storage Python objects), and subsequent operations — even a + bare ``.device`` read during ``_check_if_offload`` for a follow-up + ``push_tensor`` on the same original — would crash with + `` has no data!``. + """ + row = None + if self._rowwise_storage is not None: + row_cls = type(self._rowwise_storage) + if hasattr(row_cls, "make_like"): + row = row_cls.make_like(self._rowwise_storage) + else: + raise NotImplementedError( + "HybridQuantizedTensor.detach() does not support storage-only " + f"rowwise sub-storage {row_cls.__name__}" + ) + col = None + if self._columnwise_storage is not None: + col_cls = type(self._columnwise_storage) + if hasattr(col_cls, "make_like"): + col = col_cls.make_like(self._columnwise_storage) + else: + raise NotImplementedError( + "HybridQuantizedTensor.detach() does not support storage-only " + f"columnwise sub-storage {col_cls.__name__}" + ) + return HybridQuantizedTensor( + shape=self.shape, + dtype=self.dtype, + rowwise_storage=row, + columnwise_storage=col, + quantizer=self._quantizer, + ) + + def get_metadata(self) -> Dict[str, Any]: + return HybridQuantizedTensorStorage.get_metadata(self) + + @staticmethod + def _move_metadata_value( + value: Any, + *, + target_device: torch.device, + non_blocking: bool, + pin_memory: bool, + ) -> Any: + if isinstance(value, torch.Tensor): + value = value.to(device=target_device, non_blocking=non_blocking) + if pin_memory and target_device.type == "cpu": + value = value.pin_memory() + return value + + @classmethod + def _move_sub_storage( + cls, + sub_storage: Optional[QuantizedTensorStorage], + *, + target_device: torch.device, + non_blocking: bool, + pin_memory: bool, + ) -> Optional[QuantizedTensorStorage]: + if sub_storage is None: + return None + metadata = { + key: cls._move_metadata_value( + value, + target_device=target_device, + non_blocking=non_blocking, + pin_memory=pin_memory, + ) + for key, value in sub_storage.get_metadata().items() + } + if isinstance(sub_storage, QuantizedTensor): + metadata.update( + { + "shape": sub_storage.shape, + "dtype": sub_storage.dtype, + "requires_grad": sub_storage.requires_grad, + "device": target_device, + } + ) + return type(sub_storage)(**metadata) + + def contiguous( + self, + memory_format: torch.memory_format = torch.contiguous_format, + ) -> "HybridQuantizedTensor": + """Return a HybridQuantizedTensor with contiguous sub-storages.""" + + def _contiguous_sub( + role: str, + sub_storage: Optional[QuantizedTensorStorage], + ) -> Optional[QuantizedTensorStorage]: + if sub_storage is None: + return None + if not isinstance(sub_storage, torch.Tensor): + raise ValueError( + "HybridQuantizedTensor.contiguous does not support storage-only " + f"{role} sub-storage {type(sub_storage).__name__}. This path is " + "only supported for tensor sub-storages." + ) + try: + return sub_storage.contiguous(memory_format=memory_format) + except (NotImplementedError, ValueError) as err: + raise ValueError( + "HybridQuantizedTensor.contiguous could not make the " + f"{role} sub-storage {type(sub_storage).__name__} contiguous " + f"with memory_format={memory_format}." + ) from err + + row = _contiguous_sub("rowwise", self._rowwise_storage) + col = _contiguous_sub("columnwise", self._columnwise_storage) + if row is self._rowwise_storage and col is self._columnwise_storage: + return self + return HybridQuantizedTensor( + shape=self.shape, + dtype=self.dtype, + rowwise_storage=row, + columnwise_storage=col, + quantizer=self._quantizer, + requires_grad=self.requires_grad, + device=self.device, + ) + + def __reduce_ex__(self, protocol: int) -> tuple: + """Custom pickling. + + Without this, the default ``torch.Tensor.__reduce_ex__`` rebuilds + the parameter as a plain ``torch.Tensor``, dropping the + sub-storages and per-tensor scale state. DCP then reloads the + parameter via ``aten.copy_(dst, plain_tensor)`` which routes to + ``dst.quantize_(plain_tensor)`` — re-quantizing dequantized data + loses precision. + + Mirrors the per-format ``__reduce_ex__`` on ``Float8Tensor``, + ``MXFP8Tensor``, ``NVFP4Tensor``, and ``Float8BlockwiseQTensor``. + Each sub-storage is itself pickled via its own ``__reduce_ex__`` + (preserving FP8 bytes + ``_scale_inv``); the quantizers travel as + regular Python objects and must therefore be picklable + themselves. + """ + return ( + _make_hybrid_quantized_tensor_in_reduce_ex, + ( + self._rowwise_storage, + self._columnwise_storage, + self._quantizer, + self.dtype, + self.shape, + ), + ) + + # ── FSDP2 protocol ────────────────────────────────────────────── + + def fsdp_pre_all_gather( # pylint: disable=unused-argument + self, mesh, orig_size, contiguous_orig_stride, module, mp_policy + ): + """Extract plain tensor buffers from both sub-storages for FSDP2 all-gather. + + Always send both directions. This gives a stable buffer count/shape + across forward and backward, at the cost of gathering the unused + direction each pass. No requantization, no BF16 fallback. + + Buffer extraction is delegated to each sub-storage's + :meth:`QuantizedTensorStorage.fsdp_extract_buffers`, which strips any + format-specific padding (e.g. MXFP8 block-scale alignment) before the + gather so concatenation along dim-0 is well-defined. + + TODO(#3158): bandwidth optimization — pack both directions into a + single flat buffer sized ``max(row_bytes, col_bytes)`` (not + ``row_bytes + col_bytes``) to halve comm volume for asymmetric format + pairs. Planned implementation: a new per-sub-storage + ``fsdp_pack_into(flat_buffer, offset, meta)`` helper that layouts + both directions back-to-back with offsets stored in the metadata + tuple; ``fsdp_post_all_gather`` would slice the gathered flat buffer + using those offsets. + """ + # Mirror ``Float8Tensor.fsdp_pre_all_gather``: enable cross-shard amax + # reduction so the post-optimizer re-quantization of the sharded weight + # keeps one shared scale across shards (no-op for sub-quantizers without + # amax reduction, e.g. MXFP8). + if mesh is not None: + self._quantizer.amax_reduction_group = mesh.get_group() + self._quantizer.with_amax_reduction = True + + # Quick, targeted error for sub-storages whose FSDP2 support isn't + # implemented yet (e.g. NVFP4). Without this, users hit + # NotImplementedError from deep inside fsdp_extract_buffers with a + # generic message. + for role, sub in ( + ("rowwise", self._rowwise_storage), + ("columnwise", self._columnwise_storage), + ): + if sub is None: + continue + if not isinstance(sub, QuantizedTensor): + raise NotImplementedError( + "Hybrid FSDP2 all-gather does not support storage-only " + f"{role} sub-storage {type(sub).__name__}. This usually means " + "a HybridQuantizer sub-quantizer had internal=True; use " + "tensor sub-storages for Hybrid FSDP2 or disable Hybrid FSDP2 " + "for this parameter." + ) + try: + sub.fsdp_buffer_fields() + except NotImplementedError as err: + raise NotImplementedError( + "Hybrid FSDP2 all-gather is not supported for a " + f"{type(sub).__name__} {role} sub-storage: it does not " + "implement fsdp_buffer_fields. " + "NVFP4 sub-storages need packed-FP4 dim-0 alignment, " + "columnwise dequantization and RHT-cache handling before " + "they can be gathered. Use a supported sub-quantizer " + "(Float8CurrentScaling, MXFP8, Float8Block) or run without " + "FSDP2." + ) from err + + row_buffers: Tuple[Optional[torch.Tensor], ...] = () + col_buffers: Tuple[Optional[torch.Tensor], ...] = () + row_meta: Optional[Dict[str, Any]] = None + col_meta: Optional[Dict[str, Any]] = None + if self._rowwise_storage is not None: + row_buffers, row_meta = self._rowwise_storage.fsdp_extract_buffers() + if self._columnwise_storage is not None: + col_buffers, col_meta = self._columnwise_storage.fsdp_extract_buffers() + + sharded_tensors = row_buffers + col_buffers + + metadata = ( + len(row_buffers), + row_meta, + col_meta, + self._rowwise_storage, # original sharded sub-storage (for make_like on iter-1) + self._columnwise_storage, + self._quantizer, + ) + return sharded_tensors, metadata + + def fsdp_post_all_gather( + self, + all_gather_outputs: Tuple[torch.Tensor, ...], + metadata: Any, + param_dtype: torch.dtype, + *, + out: Optional[HybridQuantizedTensor] = None, + ): + """Reconstruct HybridQuantizedTensor from all-gathered buffers. + + On iteration 1 (``out=None``): clone each sub-storage via + :meth:`make_like` from the sharded original, then delegate the + gathered-buffer writeback (and any format-specific re-padding) to + :meth:`QuantizedTensorStorage.fsdp_assign_gathered`. + On iteration 2+ (``out=prev``): delegate directly to the existing + sub-storages' ``fsdp_assign_gathered``. + """ + ( + n_row_buffers, + row_meta, + col_meta, + orig_row_sub, + orig_col_sub, + hybrid_quantizer, + ) = metadata + + row_quantizer = hybrid_quantizer.rowwise_quantizer + col_quantizer = hybrid_quantizer.columnwise_quantizer + + row_gathered = all_gather_outputs[:n_row_buffers] + col_gathered = all_gather_outputs[n_row_buffers:] + + def _infer_shape(gathered_buffers): + for buf in gathered_buffers: + if buf is not None: + return buf.shape + return None + + # ``update_usage`` after gathered writeback mirrors what vanilla + # ``Float8Tensor.fsdp_post_all_gather`` does + # — invalidates any stale ``_transpose`` cache on Float8 sub-storages + # and recreates the transpose on non-Hopper architectures where the + # FP8 cuBLAS path requires it. No-op on Blackwell. The flags come + # from the sub-quantizer's pinned direction (set by + # ``HybridQuantizer.__init__``), so we honor whatever the inner + # quantizer thinks its direction is. + def _sync_usage(sub_storage, sub_quantizer): + if sub_storage is None or sub_quantizer is None: + return + sub_storage.update_usage( + rowwise_usage=sub_quantizer.rowwise_usage, + columnwise_usage=sub_quantizer.columnwise_usage, + ) + + if out is not None: + # Iteration 2+: in-place field update on existing sub-storages + if out._rowwise_storage is not None and row_meta is not None: + out._rowwise_storage.fsdp_assign_gathered(row_gathered, row_meta) + _sync_usage(out._rowwise_storage, out._quantizer.rowwise_quantizer) + if out._columnwise_storage is not None and col_meta is not None: + out._columnwise_storage.fsdp_assign_gathered(col_gathered, col_meta) + _sync_usage(out._columnwise_storage, out._quantizer.columnwise_quantizer) + else: + # First iteration: clone the original sharded sub-storages via make_like, + # then write gathered (full-size) buffers via each sub-storage's own + # fsdp_assign_gathered so padding is re-applied where applicable. + row_sub = None + if orig_row_sub is not None and isinstance(orig_row_sub, QuantizedTensor): + gathered_shape = _infer_shape(row_gathered) + row_sub = type(orig_row_sub).make_like(orig_row_sub, shape=gathered_shape) + if row_meta is not None: + row_sub.fsdp_assign_gathered(row_gathered, row_meta) + _sync_usage(row_sub, row_quantizer) + + col_sub = None + if orig_col_sub is not None and isinstance(orig_col_sub, QuantizedTensor): + gathered_shape = _infer_shape(col_gathered) + col_sub = type(orig_col_sub).make_like(orig_col_sub, shape=gathered_shape) + if col_meta is not None: + col_sub.fsdp_assign_gathered(col_gathered, col_meta) + _sync_usage(col_sub, col_quantizer) + + ref_sub = row_sub if row_sub is not None else col_sub + out = HybridQuantizedTensor( + shape=( + ref_sub.shape + if ref_sub is not None + else _infer_shape(row_gathered + col_gathered) + ), + dtype=param_dtype, + rowwise_storage=row_sub, + columnwise_storage=col_sub, + quantizer=hybrid_quantizer, + ) + + return out, all_gather_outputs + + @classmethod + def _delegate_reshape_op(cls, func, tensor, args, kwargs): + """Delegate a shape-altering op (slice, as_strided) to each sub-storage. + + Returns a new ``HybridQuantizedTensor`` when every non-None sub-storage + returns a ``QuantizedTensorStorage`` of the same kind (i.e. real + op support, as Float8Tensor provides via its own + ``__torch_dispatch__``). Returns ``None`` when any sub-storage + dequantized to a plain ``torch.Tensor`` (i.e. the sub-storage does not + support this op — MXFP8Tensor / Float8BlockwiseQTensor fall through + that way for real slicing today). On ``None`` the caller should defer + to ``super().__torch_dispatch__`` for a consistent BF16 fallback. + """ + + def _delegate(sub): + if sub is None: + return None + return func(sub, *args[1:], **kwargs) + + row_out = _delegate(tensor._rowwise_storage) + col_out = _delegate(tensor._columnwise_storage) + + row_ok = row_out is None or isinstance(row_out, QuantizedTensorStorage) + col_ok = col_out is None or isinstance(col_out, QuantizedTensorStorage) + if not (row_ok and col_ok): + return None + if row_out is None and col_out is None: + return None + + ref = row_out if row_out is not None else col_out + return HybridQuantizedTensor( + shape=ref.shape, + dtype=tensor.dtype, + rowwise_storage=row_out, + columnwise_storage=col_out, + quantizer=tensor._quantizer, + ) + + @staticmethod + def _new_zeroed_sub_storage( + sub_storage: QuantizedTensorStorage, + shape: torch.Size, + *, + dtype: torch.dtype, + device: torch.device, + pin_memory: bool, + ) -> QuantizedTensorStorage: + """Allocate and initialize a zero sub-storage without touching live state.""" + + try: + quantizer = sub_storage._get_quantizer().copy() + out = quantizer.make_empty( + shape, + dtype=dtype, + device=device, + requires_grad=False, + pin_memory=pin_memory, + ) + except NotImplementedError as exc: + raise NotImplementedError( + "HybridQuantizedTensor.new_zeros cannot construct zero-initialized " + f"{type(sub_storage).__name__} storage" + ) from exc + + if not isinstance(out, type(sub_storage)): + raise NotImplementedError( + "HybridQuantizedTensor.new_zeros did not preserve sub-storage format: " + f"expected {type(sub_storage).__name__}, got {type(out).__name__}" + ) + + buffers, storage = out.prepare_for_saving() + try: + with torch.no_grad(): + for buffer in buffers: + if buffer is not None: + buffer.zero_() + finally: + leftover = storage.restore_from_saved(buffers) + if leftover: + raise RuntimeError( + f"{type(out).__name__}.restore_from_saved did not consume all buffers" + ) + return out + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs=None): + if kwargs is None: + kwargs = {} + + if func == aten.detach.default: + return args[0].detach() + + if func == aten._to_copy.default: + tensor = args[0] + kw = dict(kwargs) if kwargs else {} + dtype = kw.get("dtype", None) + if dtype is None or dtype == tensor.dtype: + target_device = torch.device(kw.get("device", tensor.device) or tensor.device) + pin_memory = bool(kw.get("pin_memory", False)) + non_blocking = bool(kw.get("non_blocking", False)) + row = cls._move_sub_storage( + tensor._rowwise_storage, + target_device=target_device, + non_blocking=non_blocking, + pin_memory=pin_memory, + ) + col = cls._move_sub_storage( + tensor._columnwise_storage, + target_device=target_device, + non_blocking=non_blocking, + pin_memory=pin_memory, + ) + return HybridQuantizedTensor( + shape=tensor.shape, + dtype=tensor.dtype, + rowwise_storage=row, + columnwise_storage=col, + quantizer=tensor._quantizer, + requires_grad=tensor.requires_grad, + device=target_device, + ) + + # ── FSDP2: view ────────────────────────────────────────────── + if func == aten.view.default: + tensor = args[0] + shape = args[1] + # Identity view fast-path (FSDP2 reset_sharded_param issues a view + # to the param's own shape). The columnwise sub-storage's own shape + # is transposed relative to the hybrid for some formats (e.g. a 2D + # block-scaled Float8BlockwiseQTensor has shape (N, M) for an + # (M, N) weight). Forwarding the hybrid's row-major shape to it + # would be a spurious last-2-dims change, which 2D block scaling + # cannot represent and so dequantizes the sub-storage to a plain + # tensor -- breaking the FSDP2 sub-storage protocol later. Preserve + # the sub-storages as-is, mirroring the as_strided / slice no-op + # fast paths below. + if list(shape) == list(tensor.shape): + return HybridQuantizedTensor.make_like(tensor) + row_view = None + col_view = None + if tensor._rowwise_storage is not None: + row_view = tensor._rowwise_storage.view(shape) + if tensor._columnwise_storage is not None: + col_view = tensor._columnwise_storage.view(shape) + return HybridQuantizedTensor( + shape=shape, + dtype=tensor.dtype, + rowwise_storage=row_view, + columnwise_storage=col_view, + quantizer=tensor._quantizer, + ) + + # ── FSDP2: split ───────────────────────────────────────────── + if func == aten.split.Tensor: + tensor = args[0] + split_size = args[1] + dim = kwargs.get("dim", args[2] if len(args) > 2 else 0) + + if dim != 0: + return super().__torch_dispatch__(func, types, args, kwargs) + + row_pieces = ( + torch.split(tensor._rowwise_storage, split_size, dim=dim) + if tensor._rowwise_storage is not None + else None + ) + col_pieces = ( + torch.split(tensor._columnwise_storage, split_size, dim=dim) + if tensor._columnwise_storage is not None + else None + ) + + if row_pieces is None and col_pieces is None: + return super().__torch_dispatch__(func, types, args, kwargs) + + # TODO(#3158): Support Hybrid sub-storages that fall back to a + # high-precision tensor for an unquantizable local shard. + for direction, pieces in ( + ("rowwise", row_pieces), + ("columnwise", col_pieces), + ): + if pieces is None: + continue + for piece in pieces: + if not isinstance(piece, QuantizedTensor): + raise NotImplementedError( + "HybridQuantizedTensor split produced a high-precision " + f"{type(piece).__name__} for the {direction} sub-storage " + f"with local shape {tuple(piece.shape)}. Hybrid FSDP2 does " + "not support high-precision fallback children. MXFP8 FSDP " + "shards must have a first dimension divisible by 32; adjust " + "the sharding topology or disable Hybrid MXFP8 for this " + "parameter. See #3158." + ) + + num_pieces = len(row_pieces) if row_pieces is not None else len(col_pieces) + return [ + HybridQuantizedTensor( + shape=(row_pieces[i].shape if row_pieces is not None else col_pieces[i].shape), + dtype=tensor.dtype, + rowwise_storage=row_pieces[i] if row_pieces is not None else None, + columnwise_storage=col_pieces[i] if col_pieces is not None else None, + quantizer=tensor._quantizer, + ) + for i in range(num_pieces) + ] + + # ── FSDP2: as_strided / slice ──────────────────────────────── + # Fast path for no-op (common during FSDP2 reset_sharded_param); + # otherwise delegate per sub-storage so we inherit each sub-storage's + # own support level. Float8Tensor implements real slicing/as_strided + # via `_data.__torch_dispatch__`; MXFP8Tensor and Float8BlockwiseQTensor + # handle only the no-op case and fall through to dequantize for real + # ops (matching their vanilla FSDP2 behaviour). If any sub-storage + # returns a plain torch.Tensor (dequantized), we can't rewrap into a + # hybrid so we fall through to super() for a consistent BF16 fallback. + if func == aten.as_strided.default: + tensor = args[0] + shape = args[1] + strides = args[2] + storage_offset = kwargs.get("storage_offset", args[3] if len(args) > 3 else None) + if storage_offset is None: + storage_offset = tensor.storage_offset() + if ( + tuple(shape) == tuple(tensor.size()) + and tuple(strides) == tuple(tensor.stride()) + and storage_offset == tensor.storage_offset() + ): + return HybridQuantizedTensor.make_like(tensor) + out = cls._delegate_reshape_op(func, tensor, args, kwargs) + if out is not None: + return out + return super().__torch_dispatch__(func, types, args, kwargs) + + if func == aten.slice.Tensor: + tensor = args[0] + dim = args[1] + start = args[2] + end = args[3] + step = args[4] if len(args) > 4 else 1 + if start == 0 and end == tensor.size(dim) and step == 1: + return HybridQuantizedTensor.make_like(tensor) + out = cls._delegate_reshape_op(func, tensor, args, kwargs) + if out is not None: + return out + return super().__torch_dispatch__(func, types, args, kwargs) + + # ── FSDP2: copy_ ───────────────────────────────────────────── + # Fast path for hybrid-to-hybrid (FSDP2 fills buffer allocated via + # new_zeros/make_empty). Other src types (e.g. a BF16 master weight + # during checkpoint load) fall through to QuantizedTensor's base + # dispatch which routes to ``dst.quantize_(src)``. + if func == aten.copy_.default: + dst, src = args[0], args[1] + if isinstance(dst, HybridQuantizedTensor) and isinstance(src, HybridQuantizedTensor): + dst_usages = dst.get_usages() + src_usages = src.get_usages() + if dst_usages != src_usages: + raise NotImplementedError( + "HybridQuantizedTensor.copy_ requires matching rowwise/columnwise " + f"usages, but source has {src_usages} and destination has " + f"{dst_usages}. Copy from a high-precision tensor instead." + ) + + if dst._rowwise_storage is not None and src._rowwise_storage is not None: + aten.copy_.default(dst._rowwise_storage, src._rowwise_storage) + if dst._columnwise_storage is not None and src._columnwise_storage is not None: + aten.copy_.default(dst._columnwise_storage, src._columnwise_storage) + return dst + + # ── FSDP2: new_zeros ───────────────────────────────────────── + if func == aten.new_zeros.default: + tensor = args[0] + new_shape = torch.Size(args[1]) + if tensor._rowwise_storage is None and tensor._columnwise_storage is None: + raise RuntimeError( + "HybridQuantizedTensor.new_zeros requires at least one present sub-storage" + ) + dtype = kwargs.get("dtype") or tensor.dtype + device = torch.device(kwargs.get("device") or tensor.device) + layout = kwargs.get("layout") + pin_memory = bool(kwargs.get("pin_memory")) + if layout is not None and layout != torch.strided: + raise NotImplementedError( + "HybridQuantizedTensor.new_zeros only supports torch.strided layout" + ) + + supported_quantized_dtypes = (torch.float32, torch.float16, torch.bfloat16) + for direction, sub_storage in ( + ("rowwise", tensor._rowwise_storage), + ("columnwise", tensor._columnwise_storage), + ): + if sub_storage is None or isinstance(sub_storage, IdentityTensorStorage): + continue + if dtype not in supported_quantized_dtypes: + raise TypeError( + "HybridQuantizedTensor.new_zeros only supports float32, float16, or " + f"bfloat16 with a {direction} {type(sub_storage).__name__}; got {dtype}." + ) + + row = ( + cls._new_zeroed_sub_storage( + tensor._rowwise_storage, + new_shape, + dtype=dtype, + device=device, + pin_memory=pin_memory, + ) + if tensor._rowwise_storage is not None + else None + ) + col = ( + cls._new_zeroed_sub_storage( + tensor._columnwise_storage, + new_shape, + dtype=dtype, + device=device, + pin_memory=pin_memory, + ) + if tensor._columnwise_storage is not None + else None + ) + + # The source's parent usage flags are mutable and may no longer + # describe its surviving sub-storages. Give the result an isolated + # parent whose dynamic usage matches what was actually allocated; + # otherwise a later copy_ from a plain tensor can silently skip a + # present direction. Bind its children to copies of the newly + # allocated storage quantizers so their format/state stays coherent. + quantizer = tensor._quantizer.copy() + if row is not None: + quantizer.rowwise_quantizer = row._get_quantizer().copy() + quantizer.rowwise_quantizer.set_usage(rowwise=True, columnwise=False) + if col is not None: + quantizer.columnwise_quantizer = col._get_quantizer().copy() + quantizer.columnwise_quantizer.set_usage(rowwise=False, columnwise=True) + quantizer.set_usage( + rowwise=row is not None, + columnwise=col is not None, + ) + + return HybridQuantizedTensor( + shape=new_shape, + dtype=dtype, + rowwise_storage=row, + columnwise_storage=col, + quantizer=quantizer, + requires_grad=False, + device=device, + ) + + # ── FSDP2: clone ───────────────────────────────────────────── + if func == aten.clone.default: + tensor = args[0] + row_clone = ( + torch.clone(tensor._rowwise_storage) + if tensor._rowwise_storage is not None + else None + ) + col_clone = ( + torch.clone(tensor._columnwise_storage) + if tensor._columnwise_storage is not None + else None + ) + return HybridQuantizedTensor( + shape=tensor.shape, + dtype=tensor.dtype, + rowwise_storage=row_clone, + columnwise_storage=col_clone, + quantizer=tensor._quantizer, + ) + + return super().__torch_dispatch__(func, types, args, kwargs) + + +def _make_hybrid_quantized_tensor_in_reduce_ex( + rowwise_storage: Optional[QuantizedTensorStorage], + columnwise_storage: Optional[QuantizedTensorStorage], + quantizer: HybridQuantizer, + dtype: torch.dtype, + shape: torch.Size, +) -> HybridQuantizedTensor: + """Reconstruct a ``HybridQuantizedTensor`` from its ``__reduce_ex__`` payload.""" + return HybridQuantizedTensor( + shape=shape, + dtype=dtype, + rowwise_storage=rowwise_storage, + columnwise_storage=columnwise_storage, + quantizer=quantizer, + ) diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py new file mode 100644 index 0000000000..9fb980a755 --- /dev/null +++ b/transformer_engine/pytorch/tensor/identity_tensor.py @@ -0,0 +1,355 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""High-precision passthrough quantizer and tensor. + +``IdentityQuantizer`` stores a tensor directly without a low-precision encoding, +preserving the input dtype by default. It exists so the ``CustomRecipe`` + +``qfactory`` machinery can express *high-precision* tensors and, composed inside +a :class:`HybridQuantizer`, *high-precision directions* without scattering +``None``/``isinstance`` special-cases across the modules. Here high precision +means the held compute dtype, typically BF16, FP16, or FP32; it does not mean +FP32 specifically. +""" + +from __future__ import annotations +from typing import Any, Iterable, Optional, Tuple, Union + +import torch +from torch.ops import aten + +from .storage.identity_tensor_storage import IdentityTensorStorage +from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer + + +class IdentityQuantizer(Quantizer): + """Quantizer that produces a high-precision passthrough representation. + + Returns an :class:`IdentityTensorStorage` (or :class:`IdentityTensor`) + holding the tensor directly, without a low-precision encoding. + ``general_gemm`` materializes it as a plain tensor, so a GEMM consumes it + in the held dtype. + + Parameters + ---------- + dtype : torch.dtype, optional + If set, the held tensor is cast to this dtype on quantize. ``None`` + (default) keeps the input's dtype. + rowwise, columnwise : bool + Usage flags (kept for interface compatibility; the single + high-precision buffer serves both directions). + """ + + def __init__( + self, + *, + dtype: Optional[torch.dtype] = None, + rowwise: bool = True, + columnwise: bool = True, + ) -> None: + super().__init__(rowwise=rowwise, columnwise=columnwise) + self.dtype = dtype + + def copy(self) -> "IdentityQuantizer": + """Create shallow copy.""" + quantizer = IdentityQuantizer( + dtype=self.dtype, + rowwise=self.rowwise_usage, + columnwise=self.columnwise_usage, + ) + quantizer.internal = self.internal + quantizer.optimize_for_gemm = self.optimize_for_gemm + return quantizer + + def _maybe_cast(self, tensor: torch.Tensor) -> torch.Tensor: + # Detach so the held buffer is plain "data" with no autograd graph edge, + # mirroring the real quantizers (whose quantize kernels emit fresh, + # non-differentiable tensors). Autograd connectivity for the *quantize* + # op is provided separately by ``_QuantizeFunc`` in ``Quantizer.quantize``; + # the surrounding TE module Function computes dgrad/wgrad manually. Without + # the detach the produced tensor aliases a grad-requiring input (e.g. the + # weight workspace returned across the module Function boundary), which + # creates a spurious empty grad edge. + out = tensor.detach() + if self.dtype is not None and out.dtype != self.dtype: + return out.to(self.dtype) + return out + + def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensorStorage: + data = self._maybe_cast(tensor) + if self.internal: + return IdentityTensorStorage( + hp_data=data, + fake_dtype=data.dtype, + quantizer=self, + ) + # requires_grad=False: this is the quantized "data" tensor. Autograd + # connectivity is provided by ``_QuantizeFunc`` in ``Quantizer.quantize`` + # (mirrors the real quantizers, which return non-differentiable data). + return IdentityTensor( + data.shape, + data.dtype, + hp_data=data, + quantizer=self, + requires_grad=False, + device=data.device, + ) + + def is_requantization_safe(self) -> bool: + """Identity quantization is deterministic.""" + return True + + def make_empty( + self, + shape: Iterable[int], + *, + dtype: torch.dtype = torch.float32, + device: Optional[torch.device] = None, + requires_grad: bool = False, + pin_memory: bool = False, + ) -> Union["IdentityTensor", IdentityTensorStorage]: + if device is None: + device = torch.device("cuda") + device = torch.device(device) + data_dtype = self.dtype if self.dtype is not None else dtype + data = torch.empty(tuple(shape), dtype=data_dtype, device=device, pin_memory=pin_memory) + if self.internal: + return IdentityTensorStorage( + hp_data=data, + fake_dtype=data_dtype, + quantizer=self, + ) + return IdentityTensor( + data.shape, + data_dtype, + hp_data=data, + quantizer=self, + requires_grad=requires_grad, + device=device, + ) + + def update_quantized( + self, + src: torch.Tensor, + dst: QuantizedTensorStorage, + *, + noop_flag: Optional[torch.Tensor] = None, + ) -> QuantizedTensorStorage: + if not isinstance(dst, IdentityTensorStorage): + raise ValueError( + f"IdentityQuantizer can only update IdentityTensorStorage, got {type(dst).__name__}" + ) + data = self._maybe_cast(src) + if ( + dst._hp_data is not None + and dst._hp_data.shape == data.shape + and dst._hp_data.dtype == data.dtype + and dst._hp_data.device == data.device + ): + if noop_flag is None: + dst._hp_data.copy_(data) + else: + torch.where(noop_flag == 0, data, dst._hp_data, out=dst._hp_data) + else: + if noop_flag is not None and noop_flag.item() != 0: + return dst + dst._hp_data = data.detach() + dst._dtype = data.dtype + return dst + + def calibrate(self, tensor: torch.Tensor) -> None: + # No state to calibrate. + return + + def _get_compatible_recipe(self): + # Only reachable via CustomRecipe (qfactory returns IdentityQuantizer). + from transformer_engine.common.recipe import CustomRecipe # avoid circular import + + return CustomRecipe + + +class IdentityTensor(IdentityTensorStorage, QuantizedTensor): + """High-precision passthrough tensor produced by :class:`IdentityQuantizer`. + + Presents as a standard tensor of its nominal dtype; internally it just + holds data directly in that dtype, without a low-precision encoding. + """ + + def __repr__(self, *, tensor_contents=None): + return f"IdentityTensor(dtype={self.dtype}, data={self._hp_data})" + + def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: + return IdentityTensorStorage.dequantize(self, dtype=dtype) + + def view(self, *shape) -> "IdentityTensor": + # pylint: disable=missing-function-docstring + flat_shape = shape[0] if len(shape) == 1 and not isinstance(shape[0], int) else shape + return self._wrap_data_view(self._hp_data.view(*flat_shape)) + + def detach(self) -> "IdentityTensor": + # pylint: disable=missing-function-docstring + return self._wrap_data_view(self._hp_data.detach(), requires_grad=False) + + def clone(self) -> "IdentityTensor": + # pylint: disable=missing-function-docstring + data = self._hp_data.detach().clone() if self._hp_data is not None else None + return IdentityTensor( + self.shape, + self.dtype, + hp_data=data, + quantizer=self._quantizer, + requires_grad=self.requires_grad, + device=self.device, + stride=data.stride(), + storage_offset=data.storage_offset(), + ) + + def contiguous( + self, + memory_format: torch.memory_format = torch.contiguous_format, + ) -> "IdentityTensor": + """Return an IdentityTensor with contiguous high-precision storage.""" + if self._hp_data is not None and self._hp_data.is_contiguous(memory_format=memory_format): + return self + return self._wrap_data_view(self._hp_data.contiguous(memory_format=memory_format)) + + def __reduce_ex__(self, protocol: int) -> tuple: + """Custom pickling that preserves the high-precision payload.""" + return ( + _make_identity_tensor_in_reduce_ex, + (self._hp_data, self._quantizer, self.dtype, self.shape), + ) + + def fsdp_pre_all_gather( # pylint: disable=unused-argument + self, mesh, orig_size, contiguous_orig_stride, module, mp_policy + ): + """Extract the high-precision buffer for FSDP2 all-gather.""" + return (self._hp_data,), (self._quantizer,) + + def fsdp_post_all_gather( + self, + all_gather_outputs: Tuple[torch.Tensor, ...], + metadata: Any, + param_dtype: torch.dtype, + *, + out: Optional["IdentityTensor"] = None, + ): + """Rebuild IdentityTensor from the gathered high-precision buffer.""" + (data,) = all_gather_outputs + (quantizer,) = metadata + logical_dtype = ( + quantizer.dtype + if quantizer is not None and quantizer.dtype is not None + else param_dtype + ) + if data.dtype != logical_dtype: + raise RuntimeError( + "IdentityTensor FSDP payload dtype does not match its logical dtype: " + f"payload={data.dtype}, logical={logical_dtype}." + ) + if out is not None: + out._hp_data = data + out._dtype = logical_dtype + else: + out = IdentityTensor( + shape=data.shape, + dtype=logical_dtype, + hp_data=data, + quantizer=quantizer, + requires_grad=False, + device=data.device, + ) + return out, all_gather_outputs + + def _wrap_data_view( + self, data: torch.Tensor, *, requires_grad: Optional[bool] = None + ) -> "IdentityTensor": + requires_grad = self.requires_grad if requires_grad is None else requires_grad + return IdentityTensor( + shape=data.shape, + dtype=self.dtype, + hp_data=data, + quantizer=self._quantizer, + requires_grad=requires_grad, + device=data.device, + stride=data.stride(), + storage_offset=data.storage_offset(), + ) + + @classmethod + def _delegate_view_op(cls, func, tensor, args, kwargs): + """Apply an alias-preserving view op to the held tensor and rewrap it.""" + + result = func(tensor._hp_data, *args[1:], **kwargs) + + def _wrap(value): + if isinstance(value, torch.Tensor): + return tensor._wrap_data_view(value) + return value + + if isinstance(result, tuple): + return tuple(_wrap(value) for value in result) + if isinstance(result, list): + return [_wrap(value) for value in result] + return _wrap(result) + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs=None): + if kwargs is None: + kwargs = {} + + if func == aten.detach.default: + return args[0].detach() + + if func == aten.clone.default: + return args[0].clone() + + if func in ( + aten.view.default, + aten.split.Tensor, + aten.as_strided.default, + aten.slice.Tensor, + ): + # Preserve optional arguments exactly; omitted as_strided offset + # means reuse the input view's current storage offset, not zero. + return cls._delegate_view_op(func, args[0], args, kwargs) + + if func == aten.copy_.default: + dst, src = args[0], args[1] + if isinstance(dst, IdentityTensor): + src_data = src._hp_data if isinstance(src, IdentityTensor) else src + dst._hp_data.copy_(src_data, *args[2:], **kwargs) + return dst + + if func == aten.new_zeros.default: + tensor = args[0] + new_shape = args[1] + if tensor._quantizer is not None: + out = tensor._quantizer.make_empty( + new_shape, + dtype=kwargs.get("dtype") or tensor.dtype, + device=kwargs.get("device") or tensor.device, + pin_memory=bool(kwargs.get("pin_memory", False)), + ) + out._hp_data.zero_() + return out + + return super().__torch_dispatch__(func, types, args, kwargs) + + +def _make_identity_tensor_in_reduce_ex( + hp_data: torch.Tensor, + quantizer: Optional[Quantizer], + dtype: torch.dtype, + shape: torch.Size, +) -> IdentityTensor: + """Reconstruct an ``IdentityTensor`` from its ``__reduce_ex__`` payload.""" + return IdentityTensor( + shape=shape, + dtype=dtype, + hp_data=hp_data, + quantizer=quantizer, + requires_grad=False, + device=hp_data.device if hp_data is not None else None, + ) diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 8ad4533609..76daad8bd2 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -8,8 +8,7 @@ from __future__ import annotations from collections.abc import Iterable import math - -from typing import Optional, Tuple, Union, Any +from typing import Optional, Tuple, Union, Any, Dict import warnings import torch @@ -42,6 +41,7 @@ class MXFP8Quantizer(Quantizer): """ dtype: DType + with_2d_quantization: bool def __init__( self, @@ -49,9 +49,11 @@ def __init__( *, rowwise: bool = True, columnwise: bool = True, + with_2d_quantization: bool = False, ) -> None: super().__init__(rowwise=rowwise, columnwise=columnwise) self.dtype = DType.cast(fp8_dtype) + self.with_2d_quantization = with_2d_quantization def copy(self) -> MXFP8Quantizer: """Create shallow copy""" @@ -60,12 +62,45 @@ def copy(self) -> MXFP8Quantizer: fp8_dtype=self.dtype, rowwise=self.rowwise_usage, columnwise=self.columnwise_usage, + with_2d_quantization=self.with_2d_quantization, ) quantizer.internal = self.internal quantizer.optimize_for_gemm = self.optimize_for_gemm return quantizer + # ----- TensorSpec / pure-Python allocation ----- + + def storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + return { + "cls": MXFP8TensorStorage if self.internal else MXFP8Tensor, + "nontensor_kwargs": { + "fp8_dtype": self.dtype, + "quantizer": self, + "with_gemm_swizzled_scales": self.optimize_for_gemm, + "fake_dtype": fake_dtype, + }, + } + + def inner_tensor_specs( + self, shape: Tuple[int, ...] + ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]: + shape = tuple(shape) + specs: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {} + if self.rowwise_usage: + specs["_rowwise_data"] = (shape, torch.uint8) + specs["_rowwise_scale_inv"] = ( + tuple(self.get_scale_shape(shape, columnwise=False)), + torch.uint8, + ) + if self.columnwise_usage: + specs["_columnwise_data"] = (shape, torch.uint8) + specs["_columnwise_scale_inv"] = ( + tuple(self.get_scale_shape(shape, columnwise=True)), + torch.uint8, + ) + return specs + def update_quantized( self, src: torch.Tensor, @@ -305,6 +340,10 @@ def onnx_dequantize(self, tensor: Union[MXFP8TensorStorage, MXFP8Tensor]) -> tor def _get_compatible_recipe(self) -> Union[type[Recipe], None]: return MXFP8BlockScaling + def is_requantization_safe(self) -> bool: + """MXFP8 block scales are derived deterministically from each input.""" + return True + register_value_opaque_quantizer(MXFP8Quantizer) @@ -413,8 +452,10 @@ def detach(self) -> MXFP8Tensor: def clone(self) -> MXFP8Tensor: # pylint: disable=missing-function-docstring - assert self._rowwise_data is not None - rowwise_data = self._rowwise_data.detach().clone() + # _rowwise_data may be None for columnwise-only sub-storages (hybrid quantization) + rowwise_data = ( + self._rowwise_data.detach().clone() if self._rowwise_data is not None else None + ) columnwise_data = None if self._columnwise_data is not None: columnwise_data = self._columnwise_data.detach().clone() @@ -528,77 +569,87 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): ): return super().__torch_dispatch__(func, types, args, kwargs) - out_data = [] - for data in [tensor._rowwise_data, tensor._columnwise_data]: - func_out = ( - data.__torch_dispatch__( + def _split_data(data): + if data is None: + return None + return data.__torch_dispatch__( + func, + types, + [data] + list(args[1:]), + kwargs, + ) + + row_data_splits = _split_data(tensor._rowwise_data) + col_data_splits = _split_data(tensor._columnwise_data) + + def _split_scale_inv(scale_inv, scale_split_size, pad_multiple): + if scale_inv is None: + return None + scale_inv_out = list( + scale_inv.__torch_dispatch__( func, types, - [data] + list(args[1:]), + [scale_inv, scale_split_size] + list(args[2:]), kwargs, ) - if data is not None - else None ) - out_data.append(func_out) + for idx, split_scale_inv_out in enumerate(scale_inv_out): + current_shape = split_scale_inv_out.shape + pad_dim0 = (pad_multiple - current_shape[0] % pad_multiple) % pad_multiple + if pad_dim0 > 0: + scale_inv_out[idx] = torch.nn.functional.pad( + split_scale_inv_out, (0, 0, 0, pad_dim0) + ) + return scale_inv_out - scale_invs = [tensor._rowwise_scale_inv, tensor._columnwise_scale_inv] - split_sizes_for_scale = [split_size, split_size // MXFP8_BLOCK_SCALING_SIZE] if IS_HIP_EXTENSION: if get_device_compute_capability() == (12, 5): # gfx1250 MX pre-swizzle layout requires both dims padded to multiple of 4 - padding_multiples = [4, 4] + row_pad_multiple, col_pad_multiple = 4, 4 else: - # ROCm/HIP backend uses an unpadded scale-inv layout (see `MXFP8Quantizer.make_empty`), - # so applying the padding here would produce a per-shard scale-inv whose dim-0 - # does not match the destination scale-inv allocated for the FSDP2 local shard. - padding_multiples = [1, 1] + # ROCm/HIP backend uses an unpadded scale-inv layout (see + # `MXFP8Quantizer.make_empty`), so applying the padding here would produce a + # per-shard scale-inv whose dim-0 does not match the destination scale-inv + # allocated for the FSDP2 local shard. + row_pad_multiple, col_pad_multiple = 1, 1 else: - # Padding requirements: rowwise dim0 should be divisible by 128, columnwise dim0 should be divisible by 4 - padding_multiples = [128, 4] - for scale_inv, scale_split_size, pad_multiple in zip( - scale_invs, split_sizes_for_scale, padding_multiples - ): - scale_inv_out = ( - scale_inv.__torch_dispatch__( - func, - types, - [scale_inv, scale_split_size] + list(args[2:]), - kwargs, - ) - if scale_inv is not None - else None - ) - scale_inv_out = list(scale_inv_out) if scale_inv_out is not None else None - # Pad scale_inv_out to be a multiple of pad_multiple - if scale_inv_out is not None: - for idx, split_scale_inv_out in enumerate(scale_inv_out): - current_shape = split_scale_inv_out.shape - pad_dim0 = (pad_multiple - current_shape[0] % pad_multiple) % pad_multiple - if pad_dim0 > 0: - scale_inv_out[idx] = torch.nn.functional.pad( - split_scale_inv_out, (0, 0, 0, pad_dim0) - ) - out_data.append(scale_inv_out) + # Padding requirements: rowwise dim0 should be divisible by 128, columnwise dim0 + # should be divisible by 4 + row_pad_multiple, col_pad_multiple = 128, 4 + row_scale_splits = _split_scale_inv( + tensor._rowwise_scale_inv, + split_size, + row_pad_multiple, + ) + col_scale_splits = _split_scale_inv( + tensor._columnwise_scale_inv, + split_size // MXFP8_BLOCK_SCALING_SIZE, + col_pad_multiple, + ) + + ref_splits = row_data_splits if row_data_splits is not None else col_data_splits + num_splits = len(ref_splits) return [ MXFP8Tensor( shape=( - splitted_tensor_data[0].size() - if splitted_tensor_data[0] is not None - else splitted_tensor_data[1].size() + row_data_splits[i].size() + if row_data_splits is not None + else col_data_splits[i].size() ), dtype=tensor.dtype, - rowwise_data=splitted_tensor_data[0], - rowwise_scale_inv=splitted_tensor_data[2], - columnwise_data=splitted_tensor_data[1], - columnwise_scale_inv=splitted_tensor_data[3], + rowwise_data=row_data_splits[i] if row_data_splits is not None else None, + rowwise_scale_inv=row_scale_splits[i] if row_scale_splits is not None else None, + columnwise_data=col_data_splits[i] if col_data_splits is not None else None, + columnwise_scale_inv=( + col_scale_splits[i] if col_scale_splits is not None else None + ), quantizer=tensor._quantizer, requires_grad=False, fp8_dtype=tensor._fp8_dtype, with_gemm_swizzled_scales=False, device=tensor.device, ) - for splitted_tensor_data in zip(*out_data) + for i in range(num_splits) ] if func == torch.ops.aten.as_strided.default: @@ -689,6 +740,9 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): device=tensor.device, ) + if func == torch.ops.aten.clone.default: + return cls.clone(args[0]) + # Default case return super().__torch_dispatch__(func, types, args, kwargs) diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 23027ac448..886acba632 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -9,7 +9,7 @@ from collections.abc import Iterable import math import warnings -from typing import Dict, Optional, Tuple, Union +from typing import Any, Dict, Optional, Tuple, Union import functools import torch @@ -258,6 +258,10 @@ def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: """Quantize tensor implementation""" return tex.quantize(tensor, self) + def is_requantization_safe(self) -> bool: + """NVFP4 quantization is replay-safe unless stochastic rounding is enabled.""" + return not self.stochastic_rounding + def is_quantizable(self, inp: torch.Tensor) -> bool: """Returns whether or not given inp can be quantized""" if self.row_scaled_nvfp4: @@ -349,6 +353,55 @@ def _canonicalized_amax_reduction_group(self) -> dist_group_type: def _get_compatible_recipe(self) -> Union[type[Recipe], None]: return NVFP4BlockScaling + # ----- TensorSpec / pure-Python allocation ----- + + def storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + return { + "cls": NVFP4TensorStorage if self.internal else NVFP4Tensor, + "nontensor_kwargs": { + "fp4_dtype": self.dtype, + "quantizer": self, + "with_gemm_swizzled_scales": self.optimize_for_gemm, + "row_scaled_nvfp4": self.row_scaled_nvfp4, + "nvfp4_use_4over6": self.nvfp4_use_4over6, + "nvfp4_e4m3_max": self.nvfp4_e4m3_max, + "fake_dtype": fake_dtype, + }, + } + + def inner_tensor_specs( + self, shape: Tuple[int, ...] + ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]: + shape = tuple(shape) + specs: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {} + # FP4 data packs 2 values per byte (uint8); block scales are E4M3 stored + # as uint8; amax inner tensors are FP32 (per-row when row-scaled, else scalar). + # Order matches NVFP4TensorStorage._INNER_TENSORS (the canonical + # __tensor_flatten__ order): data + scale_inv per usage first, amax last. + # Workaround: call @staticmethods via the class, not the instance -- + # instance access breaks torch.compile guard generation (pytorch #182741). + if self.rowwise_usage: + specs["_rowwise_data"] = (type(self).convert_shape_for_fp4(shape), torch.uint8) + specs["_rowwise_scale_inv"] = ( + tuple(self.get_scale_shape(shape, columnwise=False)), + torch.uint8, + ) + if self.columnwise_usage: + specs["_columnwise_data"] = ( + type(self).convert_shape_for_fp4(type(self).get_columnwise_shape(shape)), + torch.uint8, + ) + specs["_columnwise_scale_inv"] = ( + tuple(self.get_scale_shape(shape, columnwise=True)), + torch.uint8, + ) + if self.rowwise_usage: + amax_rowwise_shape = (math.prod(shape[:-1]),) if self.row_scaled_nvfp4 else (1,) + specs["_amax_rowwise"] = (amax_rowwise_shape, torch.float32) + if self.columnwise_usage: + specs["_amax_columnwise"] = ((1,), torch.float32) + return specs + register_value_opaque_quantizer(NVFP4Quantizer) @@ -677,7 +730,9 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): # View op if func == aten.view.default: if len(args) != 2: - raise RuntimeError("Unexpected args for view op (expected 2 args, got {len(args)})") + raise RuntimeError( + f"Unexpected args for view op (expected 2 args, got {len(args)})" + ) tensor = args[0] shape = args[1] if shape == list(tensor.size()): diff --git a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py index b24a4e9144..b24ea43357 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -5,18 +5,19 @@ """Mixin class holding data specific for Float8BlockwiseQTensor""" from __future__ import annotations +from collections.abc import Iterable import math -from typing import Optional, Dict, Any, Tuple, Union +from typing import Annotated, Optional, Dict, Any, Tuple, Union import torch import transformer_engine_torch as tex -from ...quantized_tensor import QuantizedTensorStorage, Quantizer +from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer from .._quantization_helpers import safe_quantized_repr from ...constants import TE_DType_To_Torch, DType -from ...utils import _empty_tensor +from ...utils import _empty_tensor, round_up_to_nearest_multiple class _FromFloat8BlockwiseFunc(torch.autograd.Function): @@ -118,12 +119,12 @@ class Float8BlockwiseQTensorStorage(QuantizedTensorStorage): be instantiated directly for performance-critical internal usage. """ - _rowwise_data: Optional[torch.Tensor] - _columnwise_data: Optional[torch.Tensor] + _rowwise_data: Annotated[Optional[torch.Tensor], InnerTensor("rowwise_data")] + _rowwise_scale_inv: Annotated[Optional[torch.Tensor], InnerTensor("rowwise_scale_inv")] + _columnwise_data: Annotated[Optional[torch.Tensor], InnerTensor("columnwise_data")] + _columnwise_scale_inv: Annotated[Optional[torch.Tensor], InnerTensor("columnwise_scale_inv")] _quantizer: Quantizer _fp8_dtype: DType - _rowwise_scale_inv: Optional[torch.Tensor] - _columnwise_scale_inv: Optional[torch.Tensor] _is_2D_scaled: bool def __new__( @@ -316,12 +317,70 @@ def size(self, *args, **kwargs): # pylint: disable=missing-function-docstring if self._rowwise_data is not None: return self._rowwise_data.size(*args, **kwargs) - dims = list(self._columnwise_data.size(*args, **kwargs)) - reordered = [] - for i in range(1, len(dims)): - reordered.append(dims[i]) - reordered.append(dims[0]) - return torch.Size(reordered) + # Columnwise data is stored transposed, so a dim argument cannot be + # forwarded to it: rebuild the logical shape first, then index into it. + dims = self._columnwise_data.shape + if len(dims) == 2: + shape = torch.Size((dims[1], dims[0])) + else: + shape = torch.Size(tuple(dims[1:]) + (dims[0],)) + dim = args[0] if args else kwargs.get("dim") + return shape if dim is None else shape[dim] + + def view(self, shape): + """Reshape the leading (token) dims without dequantizing. + + Mirrors ``MXFP8TensorStorage.view``. Only leading dims may change: block tiling fixes + the inner dim (1D) or inner two (2D), so scale-inv stays valid. Columnwise data is + stored transposed, viewed as ``[inner, *leading]``. + """ + cur_shape = self.size() + if shape is None or shape == cur_shape: + return self + # Canonicalize shape + if not isinstance(shape, Iterable): + shape = [shape] + elif len(shape) == 1 and isinstance(shape[0], Iterable): + shape = shape[0] + shape = list(shape) + if -1 in shape: + d_inferred = -math.prod(cur_shape) // math.prod(shape) + for i, d in enumerate(shape): + if d == -1: + shape[i] = d_inferred + break + if shape == list(cur_shape): + return self + + if self._is_2D_scaled: + if shape[-2:] != list(cur_shape)[-2:]: + raise RuntimeError( + "Float8BlockwiseQTensorStorage (2D block scaling) cannot reshape the inner " + f"two dimensions (attempted {tuple(cur_shape)} -> {tuple(shape)})" + ) + elif shape[-1] != cur_shape[-1]: + raise RuntimeError( + "Float8BlockwiseQTensorStorage (1D block scaling) cannot reshape the inner " + f"dimension (attempted {tuple(cur_shape)} -> {tuple(shape)})" + ) + + new_rowwise_data = None + if self._rowwise_data is not None: + new_rowwise_data = self._rowwise_data.view(*shape) + new_columnwise_data = None + if self._columnwise_data is not None: + new_columnwise_data = self._columnwise_data.view([shape[-1], *shape[:-1]]) + + return Float8BlockwiseQTensorStorage( + new_rowwise_data, + self._rowwise_scale_inv, + new_columnwise_data, + self._columnwise_scale_inv, + self._fp8_dtype, + self._quantizer, + self._is_2D_scaled, + fake_dtype=self._dtype, + ) @property def device(self): @@ -455,3 +514,154 @@ def get_usages(self) -> Dict[str, bool]: "rowwise": self._rowwise_data is not None, "columnwise": self._columnwise_data is not None, } + + # ── FSDP2 sub-storage buffer protocol ──────────────────────────── + # + # Float8Block stores columnwise data N-major (transposed) for the GEMM, so + # it cannot be dim-0 all-gathered directly. Each direction is made + # self-contained: the columnwise direction fp8-transposes its own data to + # M-major for the gather and back on assign, using only its own buffers (no + # dependency on a rowwise sibling, which in a hybrid tensor may be a + # different format). Block-scale GEMM alignment padding (round-up-to-4) is + # stripped before the gather and re-applied after. Only 2D block scaling is + # supported -- the 1D scale layout has M in dim1, incompatible with FSDP2's + # dim-0 all-gather. + + _FSDP_BLOCK_LEN = 128 + + def _fsdp_logical_mn(self) -> Tuple[int, int]: + """Flattened ``(M, N)`` of this sub-storage's logical shape.""" + shape = self.size() + last_dim = shape[-1] if len(shape) > 0 else 1 + leading = 1 + for dim in shape[:-1]: + leading *= dim + return leading, last_dim + + def fsdp_buffer_fields(self) -> Tuple[str, ...]: + """Fields gathered by FSDP2 for Float8 block scaling (2D scaling only).""" + if not self._is_2D_scaled: + raise NotImplementedError( + "FSDP2 for Float8BlockwiseQTensor requires 2D block scaling " + "(block_scaling_dim=2). 1D block scaling is not supported because " + "its scale layout has M in dim1, which is incompatible with FSDP2 " + "dim-0 all-gather." + ) + fields = [] + if self._rowwise_data is not None: + fields.extend(("_rowwise_data", "_rowwise_scale_inv")) + if self._columnwise_data is not None: + fields.extend(("_columnwise_data", "_columnwise_scale_inv")) + return tuple(fields) + + def fsdp_extract_buffers( + self, + ) -> Tuple[Tuple[Optional[torch.Tensor], ...], Dict[str, Any]]: + """Extract M-major, alignment-stripped buffers for dim-0 all-gather. + + Rowwise data is already M-major; columnwise data is N-major and is + fp8-transposed to M-major here (and transposed back in + :meth:`fsdp_assign_gathered`). The block-scale round-up-to-4 alignment + padding is stripped so dim-0 concatenation across shards is well-defined. + """ + names = self.fsdp_buffer_fields() + block_len = self._FSDP_BLOCK_LEN + m, n = self._fsdp_logical_mn() + if m % block_len != 0: + raise RuntimeError( + "FSDP2 cannot all-gather a 2-D Float8BlockwiseQTensor whose " + f"local flattened M dimension ({m}) is not a multiple of {block_len}; " + "the shard boundary splits a block-scale tile. Choose aligned shard " + "boundaries or use a supported non-blockwise recipe." + ) + m_tiles = (m + block_len - 1) // block_len + last_tiles = (n + block_len - 1) // block_len + + if self._rowwise_data is not None: + # Rowwise scale is (m_tiles, round_up(last_tiles, 4)); m_tiles sits in + # dim-0 (sharded/gathered) unpadded, the round-up padding is on dim-1 + # (not sharded). Strip dim-1 to the compact tile count. + scale = self._rowwise_scale_inv + if scale is not None and scale.size(1) > last_tiles: + scale = scale[:, :last_tiles].contiguous() + buffers = (self._rowwise_data, scale) + direction = "rowwise" + else: + # Columnwise data is N-major (N, M); transpose to M-major (M, N). + col_data = self._columnwise_data + if not col_data.is_contiguous(): + col_data = col_data.contiguous() + data_m = tex.fp8_transpose(col_data, self._fp8_dtype, out=None) + # Columnwise scale is (last_tiles, round_up(m_tiles, 4)); transpose to + # (round_up(m_tiles, 4), last_tiles) and strip dim-0 to m_tiles so the + # gathered (dim-0) axis is the M-tiles, matching the rowwise layout. + scale = self._columnwise_scale_inv.transpose(0, 1).contiguous() + if scale.size(0) > m_tiles: + scale = scale[:m_tiles].contiguous() + buffers = (data_m, scale) + direction = "columnwise" + + return buffers, {"direction": direction, "field_names": names} + + def fsdp_assign_gathered( + self, + gathered: Tuple[Optional[torch.Tensor], ...], + meta: Dict[str, Any], + ) -> None: + """Write gathered buffers back, re-applying transpose + scale padding. + + Inverse of :meth:`fsdp_extract_buffers`: rowwise re-pads the scale's + last-dim alignment; columnwise transposes the M-major gathered data back + to N-major and re-pads/transposes the scale to the GEMM scale layout + produced by ``get_scale_shape(..., columnwise=True)``. + """ + block_len = self._FSDP_BLOCK_LEN + direction = meta["direction"] + data, scale = gathered + if direction not in ("rowwise", "columnwise"): + raise RuntimeError(f"Invalid Float8Block FSDP gather direction: {direction!r}") + if data is None or scale is None: + raise RuntimeError( + "Float8Block FSDP gathered data and scale buffers must both be present." + ) + + m_full = 1 + for dim in data.shape[:-1]: + m_full *= dim + last_dim = data.size(-1) if data.dim() > 0 else 1 + expected_scale_shape = ( + (m_full + block_len - 1) // block_len, + (last_dim + block_len - 1) // block_len, + ) + if tuple(scale.shape) != expected_scale_shape: + raise RuntimeError( + "Float8Block FSDP gathered scale geometry does not match gathered data: " + f"got scale shape {tuple(scale.shape)}, expected {expected_scale_shape} " + f"for gathered data shape {tuple(data.shape)}." + ) + + if direction == "rowwise": + last_dim = data.size(-1) + last_tiles = (last_dim + block_len - 1) // block_len + if scale is not None: + pad = round_up_to_nearest_multiple(last_tiles, 4) - last_tiles + if pad > 0: + scale = torch.nn.functional.pad(scale, (0, pad)) + self._rowwise_data = data + self._rowwise_scale_inv = scale + return + + # Columnwise: gathered data is M-major (M_full, N); transpose to N-major. + data_m = data if data.is_contiguous() else data.contiguous() + self._columnwise_data = tex.fp8_transpose(data_m, self._fp8_dtype, out=None) + m_full = 1 + for dim in data.shape[:-1]: + m_full *= dim + m_tiles_full = (m_full + block_len - 1) // block_len + # Gathered scale is compact (m_tiles_full, last_tiles); transpose to + # (last_tiles, m_tiles_full) and re-pad the M-tile dim to multiple of 4. + scale_t = scale.transpose(0, 1).contiguous() + pad = round_up_to_nearest_multiple(m_tiles_full, 4) - m_tiles_full + if pad > 0: + scale_t = torch.nn.functional.pad(scale_t, (0, pad)) + self._columnwise_scale_inv = scale_t.contiguous() diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index 374d0e1e72..16bba391c4 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -5,14 +5,13 @@ """Mixin class holding data specific for Float8Tensor""" from __future__ import annotations -import math -from typing import Any, Dict, Optional, Tuple, Union +from typing import Annotated, Any, Dict, Optional, Tuple, Union import torch import transformer_engine_torch as tex -from ...quantized_tensor import QuantizedTensorStorage, Quantizer -from .._quantization_helpers import safe_quantized_repr +from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer +from .._quantization_helpers import _resolve_view_shape, safe_quantized_repr from ...constants import TE_DType as torch_to_transformer_engine_dtype, TE_DType_To_Torch, DType @@ -45,7 +44,23 @@ def forward( # Cast from FP8 return tex.dequantize(tensor, te_dtype) - raise NotImplementedError("Casting back from the transpose not implemented yet!") + if tensor._transpose is not None and not tensor._transpose_invalid: + # A columnwise-only tensor stores the logical last dimension first + # in its physical FP8 buffer. Dequantize that buffer as ordinary + # FP8 data, then restore logical row-major dimension order. + transpose_tensor = Float8TensorStorage( + data=tensor._transpose, + data_transpose=None, + fp8_scale_inv=tensor._scale_inv, + fp8_dtype=tensor._fp8_dtype, + fake_dtype=tensor._dtype, + ) + output = _FromFloat8Func.forward(None, transpose_tensor, dtype) + if output.dim() > 0: + output = output.movedim(0, -1).contiguous() + return output + + raise RuntimeError("Float8TensorStorage has neither rowwise nor columnwise data") @staticmethod def backward( @@ -67,15 +82,16 @@ class Float8TensorStorage(QuantizedTensorStorage): """ - _data: Optional[torch.Tensor] + _data: Annotated[Optional[torch.Tensor], InnerTensor("data")] _quantizer: Optional[Quantizer] _fp8_dtype: DType - _scale_inv: torch.Tensor # FP8 transpose cache - _transpose: Optional[torch.Tensor] + _transpose: Annotated[Optional[torch.Tensor], InnerTensor("data_transpose")] _transpose_invalid: bool + _scale_inv: Annotated[torch.Tensor, InnerTensor("fp8_scale_inv")] + def __new__( cls, *args, @@ -179,8 +195,16 @@ def size(self, *args, **kwargs): # pylint: disable=missing-function-docstring if self._data is not None: return self._data.size(*args, **kwargs) - size = self._transpose.size(*args, **kwargs) - return torch.Size([size[-1], math.prod(size[:-1])]) + # The transpose is stored as [last, *leading], so a dim argument cannot + # be forwarded to it: rebuild the logical shape first, then index into + # it. This matches the shape property on Float8Tensor. + dims = self._transpose.shape + if len(dims) == 2: + shape = torch.Size((dims[1], dims[0])) + else: + shape = torch.Size(tuple(dims[1:]) + (dims[0],)) + dim = args[0] if args else kwargs.get("dim") + return shape if dim is None else shape[dim] @property def device(self): @@ -193,12 +217,31 @@ def device(self): def view(self, shape: torch.Size): # pylint: disable=missing-function-docstring - out_data = self._data.view(shape) + out_data = self._data.view(shape) if self._data is not None else None + if out_data is not None: + out_shape = out_data.size() + else: + out_shape = _resolve_view_shape(self.size(), shape) out_transpose = None if self._transpose_invalid else self._transpose if out_transpose is not None: - out_transpose_shape = out_transpose.size() - if out_transpose_shape[0] != shape[-1] or out_transpose_shape[1:] != shape[:-1]: + if len(out_shape) == 0: + view_shape_for_transpose = out_shape + else: + view_shape_for_transpose = torch.Size((out_shape[-1], *out_shape[:-1])) + if out_transpose.shape != view_shape_for_transpose: + if self._data is None: + raise NotImplementedError( + "Float8TensorStorage view with columnwise-only data is only " + "supported when the requested shape preserves the columnwise layout" + ) out_transpose = None + else: + out_transpose = out_transpose.view(*view_shape_for_transpose) + if self._data is None and out_transpose is None: + raise NotImplementedError( + "Float8TensorStorage view with columnwise-only data requires a valid " + "columnwise buffer" + ) return Float8TensorStorage( data=out_data, @@ -224,6 +267,10 @@ def __repr__(self): def _create_transpose(self): """Update FP8 transpose cache""" data = self._data + # Columnwise-only Float8Tensors (e.g. hybrid quantization sub-storages) + # have _data=None — nothing to transpose. + if data is None: + return if not data.is_contiguous(): data = data.contiguous() self._transpose = tex.fp8_transpose(data, self._fp8_dtype, out=self._transpose) @@ -277,3 +324,90 @@ def get_usages(self) -> Dict[str, bool]: else: usages["columnwise"] = self._transpose is not None and not self._transpose_invalid return usages + + def fsdp_buffer_fields(self) -> Tuple[str, ...]: + """Fields gathered by FSDP2 for per-tensor FP8. + + ``_scale_inv`` is a per-tensor scalar; it travels through the hook's + metadata tuple (mirroring :meth:`Float8Tensor.fsdp_pre_all_gather`). + + Direction-aware: a vanilla Float8Tensor parameter has ``_data`` + populated, but a columnwise-only sub-storage (used inside + ``HybridQuantizedTensor`` on Hopper / L40 where non-TN FP8 GEMM is + not natively supported) holds its quantized data in ``_transpose`` + instead. Returning ``("_data",)`` unconditionally would have + ``fsdp_extract_buffers`` produce ``(None,)`` and FSDP2 would + all-gather a ``None`` tensor. + + The per-sub-storage direction is fixed at construction (pinned by + ``HybridQuantizer.__init__`` via ``set_usage``), so this check is + stable across iterations even though it inspects the current + field state. + """ + if self._data is not None: + return ("_data",) + if self._transpose is not None: + return ("_transpose",) + # Degenerate: fully empty storage. Fall back to ``_data`` so the + # base ``fsdp_extract_buffers`` returns ``(None,)`` — same surface + # the caller would have seen pre-direction-aware logic. + return ("_data",) + + def fsdp_extract_buffers( + self, + ) -> Tuple[Tuple[Optional[torch.Tensor], ...], Dict[str, Any]]: + """Extract dim-0 gather-safe buffers for per-tensor Float8. + + Rowwise ``_data`` is already M-major. A transpose-only Hopper/L40 + payload is physically ``[N, *M]`` and cannot be gathered directly + because FSDP2 grows physical dimension 0. Move N back to the logical + last dimension and make the transport buffer contiguous so FSDP2 sees + ``[*M, N]``. :meth:`fsdp_assign_gathered` restores the persistent + columnwise layout after the collective. + """ + names = self.fsdp_buffer_fields() + if names == ("_transpose",): + if self._transpose is None or self._transpose_invalid: + raise RuntimeError( + "Float8TensorStorage cannot extract an invalid columnwise transpose" + ) + transport = self._transpose + if transport.dim() > 0: + transport = transport.movedim(0, -1).contiguous() + return (transport,), { + "field_names": names, + "transport_layout": "columnwise_m_major", + } + buffers = tuple(getattr(self, name) for name in names) + return buffers, {"field_names": names, "transport_layout": "native"} + + def fsdp_assign_gathered( + self, + gathered: Tuple[Optional[torch.Tensor], ...], + meta: Dict[str, Any], + ) -> None: + """Restore gathered Float8 buffers to their persistent layout.""" + transport_layout = meta.get("transport_layout", "native") + if transport_layout == "columnwise_m_major": + if meta.get("field_names") != ("_transpose",) or len(gathered) != 1: + raise RuntimeError( + "Float8TensorStorage got invalid metadata for columnwise FSDP transport" + ) + (transport,) = gathered + if transport is None: + raise RuntimeError( + "Float8TensorStorage got a None columnwise FSDP transport buffer" + ) + self._data = None + if transport.dim() > 0: + transport = transport.movedim(-1, 0).contiguous() + self._transpose = transport + self._transpose_invalid = False + return + if transport_layout != "native": + raise RuntimeError( + f"Float8TensorStorage got unknown FSDP transport layout {transport_layout!r}" + ) + super().fsdp_assign_gathered(gathered, meta) + if "_transpose" in meta["field_names"]: + self._transpose_invalid = False diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index 7ede75943a..3473024c03 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -654,6 +654,18 @@ def make_grouped_tensor( A GroupedTensor. """ + if quantizer is not None: + # TODO(#3158): Support Identity/Hybrid packed grouped storage. + from ..hybrid_tensor import HybridQuantizer + from ..identity_tensor import IdentityQuantizer + + if isinstance(quantizer, (IdentityQuantizer, HybridQuantizer)): + raise NotImplementedError( + "GroupedTensorStorage does not support IdentityQuantizer or " + "HybridQuantizer yet. Use separate tensors, or set " + "GroupedLinear(single_grouped_weight=False). See #3158." + ) + # Set device if device is None: device = torch.cuda.current_device() diff --git a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py new file mode 100644 index 0000000000..822bb342fe --- /dev/null +++ b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py @@ -0,0 +1,232 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Storage class for composed rowwise and columnwise representations.""" + +from __future__ import annotations +from collections.abc import Iterable +from typing import Any, Dict, Optional, Tuple + +import torch + +from ...quantized_tensor import QuantizedTensorStorage, Quantizer + + +class HybridQuantizedTensorStorage(QuantizedTensorStorage): + """Storage that composes rowwise and columnwise sub-storages. + + One sub-storage provides the rowwise representation and the other provides + the columnwise representation. Either may be absent when that direction is + not needed. The sub-storages may use different formats or include an + high-precision representation such as ``IdentityTensorStorage``. + + """ + + _rowwise_storage: Optional[QuantizedTensorStorage] + _columnwise_storage: Optional[QuantizedTensorStorage] + _quantizer: Quantizer + + def __new__( + cls, + *args, + rowwise_storage: Optional[QuantizedTensorStorage], + columnwise_storage: Optional[QuantizedTensorStorage], + quantizer: Quantizer, + fake_dtype: Optional[torch.dtype] = None, + **kwargs, + ): + if quantizer is None or not isinstance(quantizer, Quantizer): + raise TypeError( + "HybridQuantizedTensorStorage requires a parent HybridQuantizer; " + f"got {type(quantizer).__name__}." + ) + if not hasattr(quantizer, "rowwise_quantizer") or not hasattr( + quantizer, "columnwise_quantizer" + ): + raise TypeError( + "HybridQuantizedTensorStorage requires a parent HybridQuantizer " + "with rowwise_quantizer and columnwise_quantizer attributes; " + f"got {type(quantizer).__name__}." + ) + + if cls is HybridQuantizedTensorStorage: + instance = object.__new__(cls) + instance._dtype = fake_dtype if fake_dtype is not None else torch.float32 + else: + instance = super().__new__(cls, *args, fake_dtype=fake_dtype, **kwargs) + + instance._rowwise_storage = rowwise_storage + instance._columnwise_storage = columnwise_storage + instance._quantizer = quantizer.copy() + return instance + + @property + def rowwise_sub_storage(self) -> Optional[QuantizedTensorStorage]: + """The sub-storage providing rowwise quantized data.""" + return self._rowwise_storage + + @property + def columnwise_sub_storage(self) -> Optional[QuantizedTensorStorage]: + """The sub-storage providing columnwise quantized data.""" + return self._columnwise_storage + + def update_usage( + self, + rowwise_usage: Optional[bool] = None, + columnwise_usage: Optional[bool] = None, + ): + # A storage object cannot reconstruct a representation that has already + # been dropped. Validate every requested direction before mutating + # either one so mixed drop/enable requests are atomic. + if rowwise_usage and self._rowwise_storage is None: + raise RuntimeError( + "Requested rowwise usage, but HybridQuantizedTensorStorage " + "has no rowwise sub-storage" + ) + if columnwise_usage and self._columnwise_storage is None: + raise RuntimeError( + "Requested columnwise usage, but HybridQuantizedTensorStorage " + "has no columnwise sub-storage" + ) + if rowwise_usage is not None and not rowwise_usage: + self._rowwise_storage = None + if columnwise_usage is not None and not columnwise_usage: + self._columnwise_storage = None + + def clear(self): + """Deallocate both sub-storages' buffers. + + Delegates to each sub-storage's own ``clear()``; no-op when a + sub-storage is ``None`` (columnwise-only or rowwise-only hybrid). + + Used by ``cpu_offload_v1`` after the offloader has taken its own + reference to the extracted buffers, to release the GPU-resident + originals. + """ + if self._rowwise_storage is not None: + self._rowwise_storage.clear() + if self._columnwise_storage is not None: + self._columnwise_storage.clear() + + def get_usages(self) -> Dict[str, bool]: + return { + "rowwise": self._rowwise_storage is not None, + "columnwise": self._columnwise_storage is not None, + } + + def prepare_for_saving( + self, + ) -> Tuple[list[Optional[torch.Tensor]], HybridQuantizedTensorStorage]: + tensors = [] + if self._rowwise_storage is not None: + row_tensors, _ = self._rowwise_storage.prepare_for_saving() + tensors.extend(row_tensors) + if self._columnwise_storage is not None: + col_tensors, _ = self._columnwise_storage.prepare_for_saving() + tensors.extend(col_tensors) + return tensors, self + + def restore_from_saved( + self, tensors: list[Optional[torch.Tensor]] + ) -> list[Optional[torch.Tensor]]: + if self._rowwise_storage is not None: + tensors = self._rowwise_storage.restore_from_saved(tensors) + if self._columnwise_storage is not None: + tensors = self._columnwise_storage.restore_from_saved(tensors) + return tensors + + def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: + """Dequantize using the first available sub-storage.""" + if dtype is None: + dtype = self._dtype + # TODO: choose the best native sub-storage for dequantization, preferring # pylint: disable=fixme + # identity/high-precision and FP8-ish formats over FP4 when both exist. + if self._rowwise_storage is not None: + return self._rowwise_storage.dequantize(dtype=dtype) + if self._columnwise_storage is not None: + return self._columnwise_storage.dequantize(dtype=dtype) + raise RuntimeError("HybridQuantizedTensorStorage has no data to dequantize") + + def get_data_tensors(self): + """Return raw data tensors from both available sub-storages.""" + row_tensors = () + col_tensors = () + if self._rowwise_storage is not None: + result = self._rowwise_storage.get_data_tensors() + row_tensors = result if isinstance(result, tuple) else (result,) + if self._columnwise_storage is not None: + result = self._columnwise_storage.get_data_tensors() + col_tensors = result if isinstance(result, tuple) else (result,) + return row_tensors + col_tensors + + def size(self, *args, **kwargs): + """Return the logical size from the first available sub-storage.""" + if self._rowwise_storage is not None: + return self._rowwise_storage.size(*args, **kwargs) + if self._columnwise_storage is not None: + return self._columnwise_storage.size(*args, **kwargs) + raise RuntimeError("HybridQuantizedTensorStorage has no data") + + @property + def device(self): + """Return the device from the first available sub-storage.""" + if self._rowwise_storage is not None: + return self._rowwise_storage.device + if self._columnwise_storage is not None: + return self._columnwise_storage.device + raise RuntimeError("HybridQuantizedTensorStorage has no data") + + def view(self, *shape): + """View delegates to each sub-storage. Used by FSDP2 reset_sharded_param. + + Identity views are handled without forwarding a reshape to the + sub-storages: the columnwise sub-storage's own shape is transposed + relative to the hybrid for some formats (e.g. a 2D block-scaled + Float8BlockwiseQTensor has shape ``(N, M)`` for an ``(M, N)`` weight), + so forwarding the hybrid's row-major shape would be a spurious + last-2-dims change that dequantizes it to a plain tensor. + """ + flat_shape = shape[0] if len(shape) == 1 and isinstance(shape[0], Iterable) else shape + if list(flat_shape) == list(self.size()): + return HybridQuantizedTensorStorage( + rowwise_storage=self._rowwise_storage, + columnwise_storage=self._columnwise_storage, + quantizer=self._quantizer, + fake_dtype=self._dtype, + ) + row_view = self._rowwise_storage.view(*shape) if self._rowwise_storage is not None else None + col_view = ( + self._columnwise_storage.view(*shape) if self._columnwise_storage is not None else None + ) + return HybridQuantizedTensorStorage( + rowwise_storage=row_view, + columnwise_storage=col_view, + quantizer=self._quantizer, + fake_dtype=self._dtype, + ) + + def get_metadata(self) -> Dict[str, Any]: + """Return constructor metadata for make_like and serialization paths.""" + return { + "rowwise_storage": self._rowwise_storage, + "columnwise_storage": self._columnwise_storage, + "quantizer": self._quantizer, + "fake_dtype": self._dtype, + } + + def __repr__(self): + row_type = ( + type(self._rowwise_storage).__name__ if self._rowwise_storage is not None else "None" + ) + col_type = ( + type(self._columnwise_storage).__name__ + if self._columnwise_storage is not None + else "None" + ) + return ( + "HybridQuantizedTensorStorage(" + f"rowwise={row_type}, " + f"columnwise={col_type}, " + f"dtype={self._dtype})" + ) diff --git a/transformer_engine/pytorch/tensor/storage/identity_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/identity_tensor_storage.py new file mode 100644 index 0000000000..c1dea51097 --- /dev/null +++ b/transformer_engine/pytorch/tensor/storage/identity_tensor_storage.py @@ -0,0 +1,152 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Storage class for a high-precision IdentityTensor representation.""" + +from __future__ import annotations +from typing import Any, Dict, Optional, Tuple + +import torch + +from ...quantized_tensor import QuantizedTensorStorage, Quantizer +from ...utils import _empty_tensor + + +class IdentityTensorStorage(QuantizedTensorStorage): + """Passthrough storage that holds a high-precision tensor directly. + + Produced by :class:`IdentityQuantizer`. It implements the + ``QuantizedTensorStorage`` interface so it can flow through the same + module / GEMM / save-for-backward / FSDP machinery as the real quantized + storages, but it uses no low-precision encoding. ``general_gemm`` + materializes it as a plain tensor in the held dtype. + + The data is direction-agnostic -- the same tensor serves both the rowwise + and columnwise directions (the GEMM transposes via its layout flags), so a + single buffer is stored. This is what lets a ``HybridQuantizer`` mix one + quantized direction with one high-precision direction. + """ + + _hp_data: Optional[torch.Tensor] + + def __new__( + cls, + *args, + hp_data: Optional[torch.Tensor], + fake_dtype: Optional[torch.dtype] = None, + quantizer: Optional[Quantizer] = None, + **kwargs, + ): + if cls is IdentityTensorStorage: + instance = object.__new__(cls) + if fake_dtype is not None: + instance._dtype = fake_dtype + elif hp_data is not None: + instance._dtype = hp_data.dtype + else: + instance._dtype = torch.float32 + else: + instance = super().__new__(cls, *args, fake_dtype=fake_dtype, **kwargs) + instance._hp_data = hp_data + instance._quantizer = quantizer.copy() if quantizer is not None else None + return instance + + def clear(self): + """Deallocate the held tensor's memory.""" + if self._hp_data is not None: + self._hp_data.data = _empty_tensor() + + def copy_from_storage(self, src: QuantizedTensorStorage) -> None: + """Copy data from another IdentityTensorStorage.""" + if not isinstance(src, IdentityTensorStorage): + raise TypeError("copy_from_storage expects IdentityTensorStorage") + if self._hp_data is not None and src._hp_data is not None: + self._hp_data.copy_(src._hp_data) + + def get_metadata(self) -> Dict[str, Any]: + """Get this tensor's metadata.""" + return { + "hp_data": self._hp_data, + "quantizer": self._quantizer, + "fake_dtype": self._dtype, + } + + def prepare_for_saving( + self, + ) -> Tuple[list[Optional[torch.Tensor]], "IdentityTensorStorage"]: + """Prepare the tensor base for saving for backward.""" + tensors = [self._hp_data] + self._hp_data = None + return tensors, self + + def restore_from_saved( + self, tensors: list[Optional[torch.Tensor]] + ) -> list[Optional[torch.Tensor]]: + """Restore the held tensor from the saved tensors list.""" + self._hp_data = tensors[0] + return tensors[1:] + + def get_data_tensors(self, rowwise_data: bool = True, columnwise_data: bool = True): + """Get this tensor's data. The single HP buffer serves both directions.""" + if rowwise_data and columnwise_data: + return self._hp_data, None + if rowwise_data: + return self._hp_data + if columnwise_data: + return self._hp_data + raise ValueError("No data to get, both rowwise_data and columnwise_data are False") + + def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: + """Return the held high-precision tensor, casting when requested.""" + if self._hp_data is None: + raise RuntimeError("IdentityTensorStorage has no data to dequantize") + if dtype is None: + dtype = self._dtype + if self._hp_data.dtype != dtype: + return self._hp_data.to(dtype) + return self._hp_data + + def update_usage( + self, + rowwise_usage: Optional[bool] = None, + columnwise_usage: Optional[bool] = None, + ): + """No-op: the single high-precision buffer serves both directions.""" + # High-precision data is not direction-specific, so there is nothing + # to drop or synthesize. Honor the request only insofar as keeping the + # buffer (a request to drop both would leave no data, which is invalid). + + def get_usages(self) -> Dict[str, bool]: + """Get the usage of the tensor.""" + has_data = self._hp_data is not None + return {"rowwise": has_data, "columnwise": has_data} + + def size(self, *args, **kwargs): + # pylint: disable=missing-function-docstring + if self._hp_data is None: + raise RuntimeError("IdentityTensorStorage has no data") + return self._hp_data.size(*args, **kwargs) + + @property + def device(self): + """Return the device of the held tensor.""" + if self._hp_data is None: + raise RuntimeError("IdentityTensorStorage has no data!") + return self._hp_data.device + + def view(self, *shape): + # pylint: disable=missing-function-docstring + flat_shape = shape[0] if len(shape) == 1 and not isinstance(shape[0], int) else shape + return IdentityTensorStorage( + hp_data=self._hp_data.view(*flat_shape) if self._hp_data is not None else None, + fake_dtype=self._dtype, + quantizer=self._quantizer, + ) + + def fsdp_buffer_fields(self) -> Tuple[str, ...]: + """Field gathered by FSDP2 for the high-precision passthrough.""" + return ("_hp_data",) + + def __repr__(self): + return f"IdentityTensorStorage(dtype={self._dtype}, data={self._hp_data})" diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index 8f47230c7b..d1a178fd32 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -7,7 +7,7 @@ """Mixin class holding data specific for MXFP8Tensor""" from __future__ import annotations -from typing import Optional, Dict, Any, Tuple, Union +from typing import Annotated, Optional, Dict, Any, Tuple, Union from collections.abc import Iterable import math import torch @@ -17,10 +17,14 @@ from transformer_engine_torch import DType as TE_DType from torch.utils.cpp_extension import IS_HIP_EXTENSION -from ...quantized_tensor import QuantizedTensorStorage, Quantizer +from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer from .._quantization_helpers import safe_quantized_repr -from ...constants import TE_DType as torch_to_transformer_engine_dtype, DType +from ...constants import ( + TE_DType as torch_to_transformer_engine_dtype, + MXFP8_BLOCK_SCALING_SIZE, + DType, +) from ...utils import _empty_tensor @@ -76,14 +80,12 @@ class MXFP8TensorStorage(QuantizedTensorStorage): """ - # Row-scaled FP8 data - _rowwise_data: Optional[torch.Tensor] - # Column-scaled FP8 data - _columnwise_data: Optional[torch.Tensor] - # Scaling factors for row-scaled FP8 data - _rowwise_scale_inv: torch.Tensor - # Scaling factors for column-scaled FP8 data - _columnwise_scale_inv: torch.Tensor + # Row-scaled FP8 data and its scaling factors + _rowwise_data: Annotated[Optional[torch.Tensor], InnerTensor("rowwise_data")] + _rowwise_scale_inv: Annotated[torch.Tensor, InnerTensor("rowwise_scale_inv")] + # Column-scaled FP8 data and its scaling factors + _columnwise_data: Annotated[Optional[torch.Tensor], InnerTensor("columnwise_data")] + _columnwise_scale_inv: Annotated[torch.Tensor, InnerTensor("columnwise_scale_inv")] # Builder class for casting to MXFP8 _quantizer: Optional[Quantizer] @@ -332,3 +334,77 @@ def get_usages(self) -> Dict[str, bool]: "rowwise": self._rowwise_data is not None, "columnwise": self._columnwise_data is not None, } + + def fsdp_buffer_fields(self) -> Tuple[str, ...]: + """Fields gathered by FSDP2 for MXFP8. + + Block scales are per-block and direction-specific — each direction + gathers both its data buffer and its scale-inv buffer. ``None``-valued + directions (e.g. a columnwise-only sub-storage in hybrid quantization) + are excluded so the gather tuple only contains real tensors. + """ + fields = [] + if self._rowwise_data is not None: + fields.extend(("_rowwise_data", "_rowwise_scale_inv")) + if self._columnwise_data is not None: + fields.extend(("_columnwise_data", "_columnwise_scale_inv")) + return tuple(fields) + + def fsdp_extract_buffers( + self, + ) -> Tuple[Tuple[Optional[torch.Tensor], ...], Dict[str, Any]]: + """Extract MXFP8 buffers, unpadding block-scale alignment before gather. + + MXFP8 kernels require scale-inv tensors aligned to ``[128, 4]`` + (rowwise) and ``[4, 128]`` (columnwise). That padding is attached to + the local shard but would produce misaligned concatenation under + FSDP2's dim-0 all-gather. Strip it here and re-apply in + :meth:`fsdp_assign_gathered`. + """ + if self._with_gemm_swizzled_scales: + raise NotImplementedError( + "FSDP2 is only supported for MXFP8Tensors with compact scales" + ) + names = self.fsdp_buffer_fields() + buffers = [] + shape = self.size() + flattened_in_shape0 = math.prod(shape[:-1]) + for name in names: + t = getattr(self, name) + if name == "_rowwise_scale_inv" and t is not None: + if t.size(0) != flattened_in_shape0: + t = t[:flattened_in_shape0] + elif name == "_columnwise_scale_inv" and t is not None: + expected = math.ceil(flattened_in_shape0 / MXFP8_BLOCK_SCALING_SIZE) + if t.size(0) != expected: + t = t[:expected] + buffers.append(t) + return tuple(buffers), {"field_names": names} + + def fsdp_assign_gathered( + self, + gathered: Tuple[Optional[torch.Tensor], ...], + meta: Dict[str, Any], + ) -> None: + """Write gathered MXFP8 buffers back, re-padding block scales. + + Inverse of :meth:`fsdp_extract_buffers`: the gathered scale-inv tensors + are padded back up to ``[128, 4]`` / ``[4, 128]`` alignment before + being assigned to the storage. + """ + names = meta["field_names"] + if len(names) != len(gathered): + raise RuntimeError( + "MXFP8TensorStorage.fsdp_assign_gathered got " + f"{len(gathered)} buffers for {len(names)} fields" + ) + for name, buf in zip(names, gathered): + if buf is not None and name == "_rowwise_scale_inv": + pad = (128 - buf.size(0) % 128) % 128 + if pad > 0: + buf = torch.nn.functional.pad(buf, (0, 0, 0, pad)) + elif buf is not None and name == "_columnwise_scale_inv": + pad = (4 - buf.size(0) % 4) % 4 + if pad > 0: + buf = torch.nn.functional.pad(buf, (0, 0, 0, pad)) + setattr(self, name, buf) diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 093f682b57..659ff3437e 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -10,14 +10,14 @@ from collections.abc import Iterable import functools import math -from typing import Any, Dict, Optional, Tuple, Union +from typing import Annotated, Any, Dict, Optional, Tuple, Union import warnings import torch import transformer_engine_torch as tex -from ...quantized_tensor import QuantizedTensorStorage, Quantizer +from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer from .._quantization_helpers import safe_quantized_repr from ...constants import TE_DType as torch_to_transformer_engine_dtype, DType @@ -90,20 +90,15 @@ class NVFP4TensorStorage(QuantizedTensorStorage): """ - # Row-scaled FP4 data - _rowwise_data: Optional[torch.Tensor] - # Column-scaled FP4 data - _columnwise_data: Optional[torch.Tensor] - # Block scaling factors for row-scaled FP4 data - _rowwise_scale_inv: torch.Tensor - # Block scaling factors for column-scaled FP4 data - _columnwise_scale_inv: torch.Tensor - # Input absolute maximum value (used to compute tensor scale for - # row-scaled FP4 data) - _amax_rowwise: torch.Tensor - # Input absolute maximum value (used to compute tensor scale for - # column-scaled FP4 data) - _amax_columnwise: torch.Tensor + # Row-scaled FP4 data and its block scaling factors + _rowwise_data: Annotated[Optional[torch.Tensor], InnerTensor("rowwise_data")] + _rowwise_scale_inv: Annotated[torch.Tensor, InnerTensor("rowwise_scale_inv")] + # Column-scaled FP4 data and its block scaling factors + _columnwise_data: Annotated[Optional[torch.Tensor], InnerTensor("columnwise_data")] + _columnwise_scale_inv: Annotated[torch.Tensor, InnerTensor("columnwise_scale_inv")] + # Input absolute maximum values, used to compute the tensor scale + _amax_rowwise: Annotated[torch.Tensor, InnerTensor("amax_rowwise")] + _amax_columnwise: Annotated[torch.Tensor, InnerTensor("amax_columnwise")] # Builder class for casting to MXFP8 _quantizer: Optional[Quantizer] diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index 29d1f1e254..0947ed9f46 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -21,6 +21,9 @@ from .nvfp4_tensor import NVFP4Tensor, NVFP4Quantizer from .mxfp8_tensor import MXFP8Tensor, MXFP8Quantizer from .float8_blockwise_tensor import Float8BlockwiseQTensor, Float8BlockQuantizer +from .hybrid_tensor import HybridQuantizedTensor, HybridQuantizer +from .identity_tensor import IdentityQuantizer +from .storage.identity_tensor_storage import IdentityTensorStorage from ..optimizers.multi_tensor_apply import multi_tensor_applier from ..utils import is_non_tn_fp8_gemm_supported, is_fp8_fnuz from ..constants import NVFP4_BLOCK_SCALING_SIZE, DType @@ -67,12 +70,139 @@ def replace_raw_data(tensor: QuantizedTensor, new_raw_data: torch.Tensor): new_raw_data.detach().copy_(old_rowwise) tensor._rowwise_data = new_raw_data del old_rowwise + elif isinstance(tensor, IdentityTensorStorage): + old_raw_data = tensor._hp_data + if old_raw_data is None: + raise RuntimeError("IdentityTensorStorage has no data") + if old_raw_data.dtype != new_raw_data.dtype: + raise ValueError( + "The data types of raw data don't match: " + f"old dtype={old_raw_data.dtype}, new dtype={new_raw_data.dtype}" + ) + new_raw_data.detach().copy_(old_raw_data) + tensor._hp_data = new_raw_data + del old_raw_data elif isinstance(tensor, MXFP8Tensor): raise NotImplementedError("replace_raw_data for MXFP8Tensor is not supported yet") + elif isinstance(tensor, HybridQuantizedTensor): + # The distopt all-gather buffer routes at the rowwise sub-storage only; + # the columnwise sub-storage is refreshed each iteration via + # ``HybridQuantizer.update_quantized``. The underlying call delegates + # to the rowwise sub-storage's own ``replace_raw_data`` (which may + # raise for sub-storage types that don't implement it). + if tensor._rowwise_storage is None: + raise NotImplementedError( + "replace_raw_data for HybridQuantizedTensor without a rowwise " + "sub-storage is not supported." + ) + replace_raw_data(tensor._rowwise_storage, new_raw_data) else: raise ValueError(f"replace_raw_data for {type(tensor)} is not supported yet") +def _is_float8_transpose_only(tensor: QuantizedTensor) -> bool: + """Whether a Float8 tensor stores its live payload only in _transpose.""" + return ( + isinstance(tensor, Float8Tensor) + and tensor._data is None + and tensor._transpose is not None + and not tensor._transpose_invalid + ) + + +def _validate_flat_fragment( + model_weight: QuantizedTensor, master_weight: torch.Tensor, start_offset +): + """Validate a flat logical shard and return its exclusive end offset.""" + if start_offset is None: + raise ValueError("start_offset must not be None when master_weight is provided") + if start_offset < 0: + raise ValueError(f"start_offset must be non-negative, got {start_offset}") + end_offset = start_offset + master_weight.numel() + if end_offset > model_weight.numel(): + raise ValueError( + f"end_offset ({end_offset}) exceeds model_weight numel ({model_weight.numel()}), " + f"start_offset={start_offset}, master_weight numel={master_weight.numel()}" + ) + return end_offset + + +def _cast_master_weight_to_rowwise_fp8_bytes( + master_weight: torch.Tensor, + model_weight: Float8Tensor, + quantizer: Float8Quantizer, +) -> torch.Tensor: + """Cast a flat master shard to row-major FP8 bytes using ``quantizer`` scale state.""" + rowwise_quantizer = Float8Quantizer( + scale=quantizer.scale, + amax=quantizer.amax, + fp8_dtype=quantizer.dtype, + rowwise=True, + columnwise=False, + ) + raw = torch.empty((1, master_weight.numel()), dtype=torch.uint8, device=model_weight.device) + temp = rowwise_quantizer.create_tensor_from_data(raw, model_weight.dtype) + rowwise_quantizer.update_quantized(master_weight.reshape(1, -1), temp) + if temp._data is None: + raise RuntimeError("Expected rowwise Float8 temporary to populate _data") + return temp._data.reshape(-1) + + +def _update_transpose_only_float8_flat_fragment( + model_weight: QuantizedTensor, + master_weight: torch.Tensor, + start_offset, + quantizer: Float8Quantizer, +) -> bool: + """Update a logical flat shard in a transpose-only Float8 tensor. + + Hopper / L40 columnwise-only Float8 sub-storages keep their live FP8 + bytes in ``_transpose`` with physical shape ``[K, rows]`` for a logical + ``[rows, K]`` tensor. A row-major logical shard is not contiguous in that + storage, so cast the shard once and scatter the resulting FP8 bytes by + logical row into the transpose buffer. + """ + if not _is_float8_transpose_only(model_weight): + return False + + _validate_flat_fragment(model_weight, master_weight, start_offset) + numel = master_weight.numel() + if numel == 0: + return True + + shape = tuple(model_weight.shape) + if len(shape) == 0: + raise ValueError("Float8 scalar transpose-only flat update is not supported") + logical_cols = int(shape[-1]) + if logical_cols <= 0 or model_weight.numel() % logical_cols != 0: + raise ValueError(f"Invalid Float8 logical shape for transpose-only update: {shape}") + logical_rows = model_weight.numel() // logical_cols + + transpose = model_weight._transpose + if transpose.numel() != model_weight.numel(): + raise ValueError( + "Float8 transpose-only storage has unexpected numel: " + f"transpose={transpose.numel()}, logical={model_weight.numel()}" + ) + transpose_2d = transpose.reshape(logical_cols, logical_rows) + fp8_bytes = _cast_master_weight_to_rowwise_fp8_bytes(master_weight, model_weight, quantizer) + + remaining = numel + logical_offset = start_offset + src_offset = 0 + while remaining > 0: + row = logical_offset // logical_cols + col = logical_offset % logical_cols + n = min(remaining, logical_cols - col) + transpose_2d[col : col + n, row].copy_(fp8_bytes[src_offset : src_offset + n]) + logical_offset += n + src_offset += n + remaining -= n + + model_weight._transpose_invalid = False + return True + + def quantize_master_weights( model_weights, master_weights, @@ -113,12 +243,29 @@ def quantize_master_weights( blockwise_scaling_params = [] mxfp8_scaling_params = [] nvfp4_params = [] + identity_params = [] if fsdp_shard_model_weights is None: use_fsdp_shard_model_weights = False fsdp_shard_model_weights = [None] * len(model_weights) else: use_fsdp_shard_model_weights = True + # Validate the entire batch before clearing initialization state or + # populating any per-format work buckets. + for model_weight, master_weight, start_offset, fsdp_shard_model_weight in zip( + model_weights, master_weights, start_offsets, fsdp_shard_model_weights + ): + _validate_per_tensor_fp8_fsdp_hopper_policy( + model_weight, + fsdp_shard_model_weight, + ) + if isinstance(model_weight, HybridQuantizedTensor): + _validate_hybrid_partial_master_policy( + model_weight, + master_weight, + start_offset, + fsdp_shard_model_weight, + ) for model_weight, master_weight, start_offset, fsdp_shard_model_weight in zip( model_weights, master_weights, start_offsets, fsdp_shard_model_weights @@ -167,6 +314,20 @@ def quantize_master_weights( mxfp8_scaling_params.append( (model_weight, master_weight, start_offset, fsdp_shard_model_weight) ) + elif isinstance(quantizer, IdentityQuantizer): + identity_params.append( + (model_weight, master_weight, start_offset, fsdp_shard_model_weight) + ) + elif isinstance(quantizer, HybridQuantizer): + _route_hybrid_to_buckets( + model_weight, + master_weight, + start_offset, + fsdp_shard_model_weight, + delayed_scaling_params=delayed_scaling_params, + current_scaling_params=current_scaling_params, + identity_params=identity_params, + ) else: raise ValueError(f"quantize_master_weights for {type(quantizer)} is not supported yet") @@ -181,6 +342,8 @@ def quantize_master_weights( _cast_master_weights_to_fp8_mxfp8_scaling(mxfp8_scaling_params, *extra_args) if len(nvfp4_params) > 0: _cast_master_weights_to_nvfp4_2d(nvfp4_params, *extra_args) + if len(identity_params) > 0: + _cast_master_weights_to_identity(identity_params, *extra_args) def cast_master_weights_to_fp8( @@ -259,6 +422,10 @@ def _cast_master_weights_to_fp8_delayed_scaling( # master_weight may be smaller than model_weight because it could be distributed across # multiple ranks. So we need to create a dummy weight using the raw data from model_weight. if not use_fsdp_shard_model_weights: + if _update_transpose_only_float8_flat_fragment( + model_weight, master_weight, start_offset, quantizer + ): + continue shard_model_weight_raw = model_weight._data.view(-1)[start_offset:end_offset] shard_model_weight_fp8 = quantizer.create_tensor_from_data( shard_model_weight_raw.view(1, -1), @@ -404,13 +571,17 @@ def _cast_master_weights_to_fp8_current_scaling( # Cast master weight to FP8 end_offset = start_offset + master_weight.numel() - if not use_fsdp_shard_model_weights: - model_weight_fragment = model_weight.reshape(-1)[start_offset:end_offset] quantizer = Float8Quantizer( scale=scale, amax=torch.Tensor(), fp8_dtype=model_weight._fp8_dtype, ) + if not use_fsdp_shard_model_weights: + if _update_transpose_only_float8_flat_fragment( + model_weight, master_weight, start_offset, quantizer + ): + continue + model_weight_fragment = model_weight.reshape(-1)[start_offset:end_offset] if use_fsdp_shard_model_weights and not isinstance(model_weight_fragment, Float8Tensor): # NOTE: The fsdp shard model weight may be a unit8 tensor instead of # a float8 tensor. We should handle this situation properly. @@ -796,6 +967,53 @@ def _cast_master_weights_to_nvfp4_2d( ) +def _identity_storage_data(tensor): + if not isinstance(tensor, IdentityTensorStorage): + raise TypeError(f"Expected IdentityTensorStorage, got {type(tensor).__name__}") + if tensor._hp_data is None: + raise RuntimeError("IdentityTensorStorage has no data") + return tensor._hp_data + + +def _cast_master_weights_to_identity( + params, group, use_fsdp_shard_model_weights=False, manual_post_all_gather_processing=False +): + del group, manual_post_all_gather_processing + + for model_weight, master_weight, start_offset, model_weight_fragment in params: + if master_weight is None: + continue + if start_offset is None: + raise ValueError("start_offset must not be None when master_weight is provided") + if start_offset < 0: + raise ValueError(f"start_offset must be non-negative, got {start_offset}") + end_offset = start_offset + master_weight.numel() + if end_offset > model_weight.numel(): + raise ValueError( + f"end_offset ({end_offset}) exceeds model_weight numel ({model_weight.numel()}), " + f"start_offset={start_offset}, master_weight numel={master_weight.numel()}" + ) + + if use_fsdp_shard_model_weights: + target = model_weight_fragment + if target is None: + raise RuntimeError("FSDP shard model weight is required for Identity writeback") + if isinstance(target, IdentityTensorStorage): + target_flat = _identity_storage_data(target).reshape(-1) + else: + target_flat = target.reshape(-1) + target_slice = target_flat[: master_weight.numel()] + else: + target_slice = _identity_storage_data(model_weight).reshape(-1)[start_offset:end_offset] + + if target_slice.numel() != master_weight.numel(): + raise ValueError( + f"Identity target slice has {target_slice.numel()} elements, " + f"but master_weight has {master_weight.numel()}" + ) + target_slice.copy_(master_weight.reshape(-1)) + + def _cast_master_weights_to_fp8_mxfp8_scaling( params, group, use_fsdp_shard_model_weights=False, manual_post_all_gather_processing=False ): # pylint: disable=unused-argument @@ -935,6 +1153,252 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( ) +# --------------------------------------------------------------------------------------------- +# HybridQuantizer helpers for `quantize_master_weights` / `post_all_gather_processing`. +# +# Dispatch is per-direction: `_route_hybrid_to_buckets` iterates over both sub-storages +# of a `HybridQuantizedTensor` and routes each one independently into the per-format +# bucket matching its own sub-quantizer type. Row and col make their own decisions and +# can mix any pair of currently-supported sub-quantizers. +# +# Supported (per-tensor Float8 or Identity sub-quantizers, any direction): +# - Float8Quantizer (delayed scaling) +# - Float8CurrentScalingQuantizer (current scaling) +# - IdentityQuantizer (high-precision passthrough) +# +# Per-tensor Float8 works because `_cast_master_weights_to_fp8_{delayed,current}_scaling` +# accept any Float8Tensor (single direction is fine — each entry is one Float8Tensor +# with its own `_scale_inv` and the helper writes that one entry's `_data`). Each +# hybrid sub-storage IS a single-direction Float8Tensor, so we route them as two +# independent entries (into the same bucket for same-format, or into different +# buckets for cross-format Float8 — e.g. delayed row + current col). +# The FSDP-sharded columnwise-only representation on Hopper is the exception: +# neither per-tensor cast helper can safely flatten its transpose-only payload, +# so `_validate_per_tensor_fp8_fsdp_hopper_policy` rejects it before mutation. +# +# Identity routes to an exact copy bucket. Single-direction hybrid (only one +# sub-storage populated) routes the present direction only. Both-None hybrids +# raise ValueError. Per-block sub-quantizers still hit their per-direction TODO. +# +# Not supported (raise NotImplementedError per-direction + TODO): +# +# - MXFP8Quantizer as a hybrid sub-quantizer (any direction) +# TODO(#3158, hybrid-mxfp8-distopt): the distopt cast kernels +# (`tex.mxfp8_scaling_compute_partial_amax`, `tex.mxfp8_scaling_partial_cast`) +# are bidirectional — both rowwise and colwise outputs required — so they +# cannot ingest a single-direction hybrid sub-storage. (Unrelated to the +# regular `tex.quantize` kernel used by forward/backward, which natively +# supports single-direction output.) Unblocker: add single-direction +# variants of the two distopt kernels, then route hybrid sub-storages +# per-direction into `mxfp8_scaling_params` matching the Float8 path above. +# Also unlocks cross-format MXFP8 row + col. +# +# - NVFP4Quantizer as a hybrid sub-quantizer (any direction) +# TODO(#3158, hybrid-nvfp4-distopt): load-bearing blocker is the kernel assertion +# `return_identity || !use_2d_quantization` in +# `quantize_transpose_vector_blockwise_fp4.cu`, which rejects exactly the +# columnwise-only 2D configuration that `HybridQuantizer.__init__` produces +# for the col sub-quantizer. Blocks hybrid 2D NVFP4 weight construction at +# `quantized_model_init` time. 1D NVFP4 is unaffected. The assertion is an +# explicitly-marked unwritten code path, not an algorithmic limit (see the +# kernel author's note above the early-return guard). +# +# Secondary blocker (gated on the kernel fix): the distopt helper +# `_cast_master_weights_to_nvfp4_2d` writes only `_rowwise_data` and relies +# on per-tensor post-AG `_create_columnwise()` — for hybrid, the columnwise +# data needs to land in a SEPARATE col sub-storage, so the post-AG branch +# must be made hybrid-aware (derive `col_sub._columnwise_data` from +# `row_sub`'s gathered rowwise). +# +# - Float8BlockQuantizer as a hybrid sub-quantizer +# TODO(#3158, hybrid-fp8-blockwise): same shape as the NVFP4 secondary blocker — +# `_cast_master_weights_to_fp8_blockwise_scaling` writes only `_rowwise_data` +# with per-tensor post-AG `_create_columnwise()` that doesn't reach hybrid's +# separate col sub-storage. Unlike NVFP4, there is no kernel-level +# construction blocker (the Block FP8 kernel natively supports +# columnwise-only mode), so hybrid Block FP8 weights construct fine via the +# non-distopt FusedAdam path today; only the sharded-master distopt cast +# path is blocked. Unblocker is a Python-side hybrid-aware post-AG branch; +# no C++ work needed. +# +# --------------------------------------------------------------------------------------------- + + +def _validate_per_tensor_fp8_fsdp_hopper_policy( + model_weight, + fsdp_shard_model_weight, +): + """Reject unsupported transpose-only per-tensor FP8 FSDP updates on Hopper. + + The FSDP shard path can receive a ``Float8Tensor`` and pass it directly to + a per-tensor FP8 cast helper. On Hopper, both delayed and current scaling + use a transpose-only representation for columnwise-only tensors that the + sharded update path cannot flatten safely. Reject the whole batch before + any scale, cache, or storage mutation instead. + """ + if fsdp_shard_model_weight is None or is_non_tn_fp8_gemm_supported(): + return + + if isinstance(model_weight, HybridQuantizedTensor): + quantizer = model_weight._get_quantizer() + candidates = ( + (model_weight._rowwise_storage, quantizer.rowwise_quantizer), + (model_weight._columnwise_storage, quantizer.columnwise_quantizer), + ) + else: + candidates = ((model_weight, model_weight._get_quantizer()),) + + for storage, quantizer in candidates: + if ( + storage is not None + and isinstance(quantizer, (Float8Quantizer, Float8CurrentScalingQuantizer)) + and _is_float8_transpose_only(storage) + ): + raise NotImplementedError( + "Columnwise-only per-tensor FP8 quantization is not implemented for " + "FSDP updates on this architecture." + ) + + +def _validate_hybrid_partial_master_policy( + model_weight, + master_weight, + start_offset, + fsdp_shard_model_weight, +): + """Reject unsupported Hybrid column-source updates before distopt mutation.""" + row_sub = model_weight._rowwise_storage + col_sub = model_weight._columnwise_storage + quantizer = model_weight._get_quantizer() + + if col_sub is not None and quantizer.columnwise_source == "rowwise_dequantized": + raise NotImplementedError( + "quantize_master_weights does not support HybridQuantizer with a live " + "columnwise representation and " + "columnwise_source='rowwise_dequantized'. The column must be derived " + "after the rowwise update/all-gather. See #3158." + ) + + if master_weight is None or row_sub is None or col_sub is None: + return + + shard_has_both_directions = ( + isinstance(fsdp_shard_model_weight, HybridQuantizedTensor) + and fsdp_shard_model_weight._rowwise_storage is not None + and fsdp_shard_model_weight._columnwise_storage is not None + ) + if shard_has_both_directions: + return + + end_offset = _validate_flat_fragment(model_weight, master_weight, start_offset) + if start_offset == 0 and end_offset == model_weight.numel(): + return + + raise ValueError( + "quantize_master_weights cannot update a two-direction " + "HybridQuantizedTensor from a partial master shard when " + "columnwise_source='original': the one-payload distributed-optimizer " + "path cannot preserve an independently quantized columnwise value. " + "Provide full-master data or retain only the rowwise representation." + ) + + +def _route_hybrid_to_buckets( + model_weight, + master_weight, + start_offset, + fsdp_shard_model_weight, + *, + delayed_scaling_params, + current_scaling_params, + identity_params, +): + """Decompose a `HybridQuantizedTensor` into per-direction entries and route each + into the appropriate per-format bucket used by `quantize_master_weights`. + + Per-direction dispatch: each sub-storage routes independently based on its + own sub-quantizer type. Per-tensor Float8 sub-quantizers (delayed and/or + current scaling) are supported in any combination per direction; single- + direction hybrid (one sub-storage dropped via ``update_usage``) is also + supported. See the TODO block above this helper for the per-block-format + rejection rationale and unblocker shapes. + """ + row_sub = model_weight._rowwise_storage + col_sub = model_weight._columnwise_storage + sub_q_row = model_weight._quantizer.rowwise_quantizer + sub_q_col = model_weight._quantizer.columnwise_quantizer + + if row_sub is None and col_sub is None: + raise ValueError( + "quantize_master_weights called on HybridQuantizedTensor with both " + "rowwise and columnwise sub-storages dropped (via update_usage). " + "Nothing to cast — this is most likely a caller bug." + ) + + # Per-direction routing: each (sub_storage, sub_quantizer) pair selects its + # own bucket based on the sub-quantizer's type. Directions that have been + # dropped via ``update_usage`` are silently skipped. + for direction, sub_storage, sub_q in ( + ("rowwise", row_sub, sub_q_row), + ("columnwise", col_sub, sub_q_col), + ): + if sub_storage is None: + continue + shard_fragment = fsdp_shard_model_weight + if shard_fragment is not None and isinstance(shard_fragment, HybridQuantizedTensor): + shard_fragment = ( + shard_fragment._rowwise_storage + if direction == "rowwise" + else shard_fragment._columnwise_storage + ) + entry = (sub_storage, master_weight, start_offset, shard_fragment) + if isinstance(sub_q, Float8Quantizer): + # Delayed scaling: the per-format helper iterates entries + # independently and does a per-DP amax all-reduce across the bucket. + delayed_scaling_params.append(entry) + elif isinstance(sub_q, Float8CurrentScalingQuantizer): + current_scaling_params.append(entry) + elif isinstance(sub_q, IdentityQuantizer): + identity_params.append(entry) + elif isinstance(sub_q, MXFP8Quantizer): + # TODO(#3158, hybrid-mxfp8-distopt): the distopt cast kernels are + # bidirectional, so a single-direction hybrid sub-storage cannot be + # fed in. See top-of-file TODO block for the unblocker (single- + # direction variants of the two distopt kernels). + raise NotImplementedError( + "quantize_master_weights for HybridQuantizer with MXFP8Quantizer " + f"{direction} sub-quantizer is not supported yet. See the TODO " + "block above _route_hybrid_to_buckets for the unblocker shape." + ) + elif isinstance(sub_q, NVFP4Quantizer): + # TODO(#3158, hybrid-nvfp4-distopt): load-bearing blocker is the kernel + # assertion that rejects columnwise-only 2D NVFP4 — which is + # exactly what hybrid's col sub-quantizer pin produces. Secondary + # blocker (gated on the kernel fix) is the per-tensor post-AG + # `_create_columnwise()` not reaching hybrid's separate col + # sub-storage. See top-of-file TODO block for details. + raise NotImplementedError( + "quantize_master_weights for HybridQuantizer with NVFP4Quantizer " + f"{direction} sub-quantizer is not supported yet. See the TODO " + "block above _route_hybrid_to_buckets for details." + ) + elif isinstance(sub_q, Float8BlockQuantizer): + # Pending hybrid-fp8-blockwise work (#3158): same shape as the NVFP4 + # secondary blocker (and only that one — no kernel-level construction + # blocker for Block FP8). Python-side post-AG fix. See the + # _route_hybrid_to_buckets design note above for details. + raise NotImplementedError( + "quantize_master_weights for HybridQuantizer with Float8BlockQuantizer " + f"{direction} sub-quantizer is not supported yet. See the TODO " + "block above _route_hybrid_to_buckets for details." + ) + else: + raise NotImplementedError( + "quantize_master_weights for HybridQuantizer with " + f"{type(sub_q).__name__} {direction} sub-quantizer is not supported yet." + ) + + def post_all_gather_processing(model_weights: Union[torch.Tensor, List[torch.Tensor]]): """ Post-processing after all-gather for weights in distributed optimizer. @@ -943,6 +1407,9 @@ def post_all_gather_processing(model_weights: Union[torch.Tensor, List[torch.Ten - Plain pytorch tensor: noop. For NVFP4 tensors, uses batched multi-tensor processing to reduce CPU overhead. + + For `HybridQuantizedTensor`, recurses per-direction so each present + sub-storage runs its native post-processing. Identity sub-storages are no-op. """ if not isinstance(model_weights, list): model_weights = [model_weights] @@ -966,6 +1433,12 @@ def post_all_gather_processing(model_weights: Union[torch.Tensor, List[torch.Ten elif isinstance(model_weight, MXFP8Tensor): # MXFP8 scaling: no need to do anything. pass + elif isinstance(model_weight, IdentityTensorStorage): + pass + elif isinstance(model_weight, HybridQuantizedTensor): + for sub in (model_weight._rowwise_storage, model_weight._columnwise_storage): + if sub is not None: + post_all_gather_processing(sub) elif isinstance(model_weight, QuantizedTensor): raise ValueError(f"post_processing for {type(model_weight)} is not supported") diff --git a/transformer_engine/pytorch/triton/mhc.py b/transformer_engine/pytorch/triton/mhc.py index 987216e327..41f71966a2 100644 --- a/transformer_engine/pytorch/triton/mhc.py +++ b/transformer_engine/pytorch/triton/mhc.py @@ -5,38 +5,173 @@ """PyTorch wrapper functions for mHC (manifold Hyper-Connection) Triton kernels.""" import os +from typing import Optional import torch import triton from transformer_engine.common.triton.mhc import ( + _mhc_projection_bwd_fused_dphi, + _mhc_projection_bwd_fused_dx, _mhc_scale_fwd_fused, _mhc_scale_bwd_fused, - _mhc_expand_combine_with_bias_fwd, - _mhc_expand_combine_with_bias_bwd, _mhc_expand_combine_fwd, _mhc_expand_combine_bwd, _mhc_aggregate_fwd, _mhc_aggregate_bwd, _mhc_projection_fwd_fused, - _mhc_projection_bwd_fused, - _mhc_sinkhorn_fwd_fused, _mhc_sinkhorn_fwd_fused_recompute, - _mhc_sinkhorn_bwd_fused, _mhc_sinkhorn_bwd_fused_recompute, + _mhc_sinkhorn_fwd_fused, + _mhc_sinkhorn_bwd_fused, ) from transformer_engine.pytorch.cpp_extensions.gemm import general_gemm +ENFORCE_DETERMINISTIC = os.environ.get("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1") == "0" + + +def _support_tma(x: torch.Tensor): + # get_device_capability returns a (major, minor) tuple; TMA needs Hopper+ (major >= 9) + return torch.cuda.get_device_capability(x.device)[0] >= 9 + + +def _tma_aligned(t): + return (t.stride(0) * t.element_size()) % 16 == 0 and t.data_ptr() % 16 == 0 + -def check_deterministic(operator: str): +_tma_allocator_initialized = False + + +def _init_tma_allocator(): + # TMA descriptors require a global memory allocation. Registered once on first use. + global _tma_allocator_initialized # pylint: disable=global-statement + if _tma_allocator_initialized: + return + + def alloc_fn( + size: int, alignment: int, stream: Optional[int] + ): # pylint: disable=unused-argument + return torch.empty(size, device="cuda", dtype=torch.int8) + + triton.set_allocator(alloc_fn) + _tma_allocator_initialized = True + + +def check_deterministic(operator: str, use_split_k: bool = False): """ - Checks if the non-deterministic algorithm is allowed for the given operator. If not, raises an assertion error with instructions on how to allow it. - Since atomic add is used in this mHC implementation, it breaks the determinism guarantee due to non-associativity of floating point addition. + If the user enforces determinism (NVTE_ALLOW_NONDETERMINISTIC_ALGO=0), split-K/M/C reductions + are disallowed because they reduce via atomic add, which is non-associative in floating point. + The default (use_split_k=False) uses deterministic store/workspace reductions instead. + """ + if use_split_k: + assert not ENFORCE_DETERMINISTIC, ( + f"[{operator}]: use_split_k=True uses atomic add which violates determinism. Either set" + " use_split_k=False or unset NVTE_ALLOW_NONDETERMINISTIC_ALGO=0." + ) + + +def mhc_generate_mix_and_aggregate( + x: torch.Tensor, + phi: torch.Tensor, + alpha: torch.Tensor, + beta: torch.Tensor, + norm_weight: Optional[torch.Tensor] = None, + use_tf32: bool = True, + fused_grad_x_acc_buffer: Optional[torch.Tensor] = None, + use_split_k: bool = False, +): + """ + Generate the mix matrix H_pre, H_post, H_res and apply H_pre to x to aggregate n streams + This wraps projection, scale, sinkhorn, and aggregate operations into one function. + + To use mHC in your model: + ``` + layer_input, H_post, H_res = mhc_generate_mix_and_aggregate(x, phi, alpha, beta) + layer_output = layer(layer_input) # Attn / FFN layer + x = mhc_fused_expand_combine(layer_output, bias, H_post, x, H_res) + ``` + + This API accepts both BF16 and FP32 parameters, though the DeepSeek V4 recipe is: + - x: BF16 + - phi, alpha, beta: FP32 + + Parameters + ---------- + x : torch.Tensor, + input tensor of shape (s, b, C, n), where s is the sequence length, b is the batch size, C is the hidden dimension per hyper connection, and n is the number of hyper connections, + dtype is torch.bfloat16 or torch.float32 + Note that C is equal to the original hidden dimension divided by n. + phi : torch.Tensor + projection matrix of shape (N, nC), where N=2n+n*n (=24 for n=4), and nC is the hidden dimension after expansion (n times of C), + dtype is torch.bfloat16 or torch.float32 + alpha : torch.Tensor + scaling factor for H, of shape (3,), where + alpha[0] is applied to H[:, 0:n] for H_pre + alpha[1] is applied to H[:, n:2n] for H_post + alpha[2] is applied to H[:, 2n:2n+n*n] for H_res + dtype: torch.bfloat16 or torch.float32 + beta : torch.Tensor + bias term for H, of shape (1, 2*n+n*n), where + beta[0, 0:n] is applied to H[:, 0:n] for H_pre + beta[0, n:2n] is applied to H[:, n:2n] for H_post + beta[0, 2n:2n+n*n] is applied to H[:, 2n:2n+n*n] for H_res + dtype is torch.bfloat16 or torch.float32 + norm_weight : torch.Tensor or None + optional, the weight for RMSNorm, of shape (K,), which is the learnable per-element affine parameters (gamma) applied to RMSNorm + dtype is torch.bfloat16 or torch.float32 + use_tf32 : bool + whether to use TF32 for matrix multiplications + fused_grad_x_acc_buffer : Optional[torch.Tensor] + A pre-allocated buffer for inplace gradient accumulation to avoid PyTorch autograd overhead. + If not None, triton kernels will accumulate the gradient of x into this same buffer to avoid copying the gradient by PyTorch, which should be reused + during the backward of mhc_fused_aggregate, mhc_fused_expand_combine and mhc_fused_projection operations + Note: the buffer must have dtype float32, and it will be cast to the activation's dtype and be returned in mhc_fused_projection + use_split_k : bool + whether to use split-K reduction with atomic adds in the projection. Faster for large K, but + non-deterministic; requires NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 + + Returns + ------- + out : torch.Tensor + out of shape (s, b, C), which is the aggregated result after applying H_pre to x, which will be fed into attention / FFN + with the same dtype as x + H_post : torch.Tensor + H_post of shape (s, b, n), which will be used in the post-processing after attention / FFN in `mhc_fused_expand_combine` + with dtype float32 + H_res : torch.Tensor + H_res of shape (s, b, n, n), which will be used to mix the residual connection in `mhc_fused_expand_combine` + with dtype float32 """ - allow_nondeterministic = os.environ.get("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1") == "1" - assert allow_nondeterministic, ( - f"[{operator}]: This operation uses atomic add which violates determinism. Set" - " NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 to allow this non-deterministic behavior." + check_deterministic("mhc_generate_mix_and_aggregate", use_split_k) + s, b, C, n = x.shape + assert ( + n == 4 + ), "Only n=4 is supported in this implementation, where n is the Hyper Connection number" + if fused_grad_x_acc_buffer is not None: + assert ( + fused_grad_x_acc_buffer.dtype == torch.float32 + ), "fused_grad_x_acc_buffer must be fp32" + assert ( + fused_grad_x_acc_buffer.numel() == x.numel() + ), "fused_grad_x_acc_buffer.numel() must match x.numel()" + nC = n * C + H, ms = mhc_fused_projection( + x.view(s * b, nC), + phi, + norm_weight=norm_weight, + use_tf32=use_tf32, + fused_grad_x_acc_buffer=fused_grad_x_acc_buffer, + use_split_k=use_split_k, + ) + H_pre, H_post, H_res = mhc_fused_scale(H, alpha, beta, ms, n) + H_res = mhc_fused_sinkhorn(H_res.view(s, b, n, n), n, recompute_hist=True, iters=20) + out = mhc_fused_aggregate( + x, + H_pre.view(s, b, n), + n, + use_tf32=use_tf32, + fused_grad_x_acc_buffer=fused_grad_x_acc_buffer, ) + return out, H_post.view(s, b, n), H_res def mhc_fused_sinkhorn( @@ -52,6 +187,7 @@ def mhc_fused_sinkhorn( ---------- H_res : torch.Tensor input H_res matrix of shape (s, b, n, n) that needs to be normalized into a doubly stochastic matrix. + dtype is torch.bfloat16 or torch.float32 n : int number of hyper connections, where only n=4 is supported in the current implementation recompute_hist : bool @@ -63,6 +199,7 @@ def mhc_fused_sinkhorn( ------- out : torch.Tensor out of shape (s, b, n, n), which is the final H_res after Sinkhorn normalization + with the same dtype as H_res """ assert n == 4, "Only n=4 is supported in this implementation" out = mHCSinkhornOp.apply(H_res, n, recompute_hist, iters) @@ -70,7 +207,11 @@ def mhc_fused_sinkhorn( def mhc_fused_scale( - H: torch.Tensor, alpha: torch.Tensor, beta: torch.Tensor, ms: torch.Tensor, n: int + H: torch.Tensor, + alpha: torch.Tensor, + beta: torch.Tensor, + ms: torch.Tensor, + n: int, ): """ Fused scale operation to compute the scaled H matrices (see eq. 16-18, section 4.3.1 of the DeepSeek mHC paper): @@ -96,6 +237,7 @@ def mhc_fused_scale( beta[0, 0:n] is applied to H[:, 0:n] for H_pre beta[0, n:2n] is applied to H[:, n:2n] for H_post beta[0, 2n:2n+n*n] is applied to H[:, 2n:2n+n*n] for H_res + Note: we assume alpha and beta have the same dtype, and according to the DeepSeek paper they should be fp32 ms : torch.Tensor mean square for each row of H from the projection kernel, of shape (M,), used for RMSNorm scaling n : int @@ -104,15 +246,17 @@ def mhc_fused_scale( Returns ------- h_pre : torch.Tensor - Scaled H_pre of shape (M, n), which aggregates (s, b, C, n) input of a Hyper Connection block into (s, b, n) as the input of attention / MLP + Scaled H_pre of shape (M, n), which aggregates (s, b, C, n) input of a Hyper Connection block into (s, b, n) as the input of attention / MLP, + with the same dtype as H h_post : torch.Tensor - Scaled H_post of shape (M, n), which expands the output of attention / MLP of shape (s, b, n) back to (s, b, C, n) for the residual connection + Scaled H_post of shape (M, n), which expands the output of attention / MLP of shape (s, b, n) back to (s, b, C, n) for the residual connection, + with the same dtype as H h_res : torch.Tensor - Scaled H_res of shape (M, n*n), which mixes the n streams of the (s, b, C, n) input of a Hyper Connection block + Scaled H_res of shape (M, n*n), which mixes the n streams of the (s, b, C, n) input of a Hyper Connection block, + with the same dtype as H """ assert n == 4, "Only n=4 is supported in this implementation" - check_deterministic("mhc_fused_scale") out = mHCScaleFusedOp.apply(H, alpha, beta, ms, n) h_pre = out[..., :n] h_post = out[..., n : 2 * n] @@ -120,7 +264,13 @@ def mhc_fused_scale( return h_pre, h_post, h_res -def mhc_fused_aggregate(x: torch.Tensor, H_pre: torch.Tensor, n: int, use_tf32: bool = True): +def mhc_fused_aggregate( + x: torch.Tensor, + H_pre: torch.Tensor, + n: int, + use_tf32: bool = True, + fused_grad_x_acc_buffer: Optional[torch.Tensor] = None, +): """ Aggregate operation to merge n activation streams into one (see section 4.3.1 of the DeepSeek mHC paper): out = x @ H_pre: (s, b, C, n) @ (s, b, n, 1) -> (s, b, C, 1) -> (s, b, C) after squeezing the last dimension @@ -130,22 +280,36 @@ def mhc_fused_aggregate(x: torch.Tensor, H_pre: torch.Tensor, n: int, use_tf32: x : torch.Tensor input activation tensor of shape (s, b, C, n), where s is the sequence length, b is the batch size, C is the hidden dimension per hyper connection, and n is the number of hyper connections. Note that C is equal to the original hidden dimension divided by n. + dtype is torch.bfloat16 or torch.float32 H_pre: torch.Tensor input H_pre matrix of shape (s, b, n) + dtype is torch.bfloat16 or torch.float32 n: int number of hyper connections, where only n=4 is supported in the current implementation use_tf32: bool whether to use TF32 precision for matmul operations. If False, it will use ieee for better precision. This is mainly used by our unittests since TF32 precision will introduce some errors and cause tests to fail + fused_grad_x_acc_buffer : Optional[torch.Tensor] + A pre-allocated buffer for inplace gradient accumulation to avoid PyTorch autograd overhead. + If not None, triton kernels will accumulate the gradient of x into this same buffer to avoid copying the gradient by PyTorch. + This optimization requires the operation order to be mhc_fused_projection -> mhc_fused_aggregate -> mhc_fused_expand_combine. + Note: the buffer must have dtype float32, and it will be cast to the activation's dtype and be returned in mhc_fused_projection Returns ------- out: torch.Tensor - output activation tensor of shape (s, b, C), which is the aggregated output after merging n hyper connections + output activation tensor of shape (s, b, C), which is the aggregated output after merging n hyper connections, + with the same dtype as x """ assert n == 4, "Only n=4 is supported in this implementation" - check_deterministic("mhc_fused_aggregate") - out = mHCAggregateOp.apply(x, H_pre, n, use_tf32) + if fused_grad_x_acc_buffer is not None: + assert ( + fused_grad_x_acc_buffer.dtype == torch.float32 + ), "fused_grad_x_acc_buffer must be fp32" + assert ( + fused_grad_x_acc_buffer.numel() == x.numel() + ), "fused_grad_x_acc_buffer.numel() must match x.numel()" + out = mHCAggregateOp.apply(x, H_pre, n, use_tf32, fused_grad_x_acc_buffer) return out @@ -157,6 +321,7 @@ def mhc_fused_expand_combine( H_res: torch.Tensor, n: int, use_tf32: bool = True, + fused_grad_x_acc_buffer: Optional[torch.Tensor] = None, ): """ Expand and combine operation for merging n hyper connections (see section 4.3.1 of the DeepSeek mHC paper): @@ -167,27 +332,45 @@ def mhc_fused_expand_combine( ---------- f : torch.Tensor input activation tensor of shape (s, b, C), which is the output from the attention / FFN sub-layer in a transformer block + dtype is torch.bfloat16 or torch.float32 bias : torch.Tensor or None optional bias tensor of shape (C,) from the last linear layer, where f + bias is fused in this kernel for better performance + dtype is torch.bfloat16 or torch.float32 H_post : torch.Tensor input H_post matrix of shape (s, b, n) + dtype is torch.bfloat16 or torch.float32 x : torch.Tensor input activation tensor of shape (s, b, C, n), which is the hyper connection input before the aggregation operation + dtype is torch.bfloat16 or torch.float32 H_res : torch.Tensor input H_res matrix of shape (s, b, n, n) + dtype is torch.bfloat16 or torch.float32 n : int - number of hyper connections + number of hyper connections, where only n=4 is supported in the current implementation use_tf32 : bool - whether to use TF32 precision for matmul operations. If False, it will use ieee for better precision. + whether to use TF32 precision for matmul operations. If False, it will use IEEE or TF32x3 for better precision. + Due to a triton bug (https://github.com/triton-lang/triton/issues/10176), we will use TF32x3 if x is bf16 and phi is fp32, or use IEEE if otherwise. This is mainly used by our unittests since TF32 precision will introduce some errors and cause tests to fail + fused_grad_x_acc_buffer : Optional[torch.Tensor] + A pre-allocated buffer for inplace gradient accumulation to avoid PyTorch autograd overhead. + If not None, triton kernels will accumulate the gradient of x into this same buffer to avoid copying the gradient by PyTorch. + This optimization requires the operation order to be mhc_fused_projection -> mhc_fused_aggregate -> mhc_fused_expand_combine. + Note: the buffer must have dtype float32, and it will be cast to the activation's dtype and be returned in mhc_fused_projection Returns ------- out : torch.Tensor - out of shape (s, b, C, n), which is the expanded and combined output after merging n hyper connections + out of shape (s, b, C, n), which is the expanded and combined output after merging n hyper connections, + with the same dtype as x """ assert n == 4, "Only n=4 is supported in this implementation" - check_deterministic("mhc_fused_expand_combine") + if fused_grad_x_acc_buffer is not None: + assert ( + fused_grad_x_acc_buffer.dtype == torch.float32 + ), "fused_grad_x_acc_buffer must be fp32" + assert ( + fused_grad_x_acc_buffer.numel() == x.numel() + ), "fused_grad_x_acc_buffer.numel() must match x.numel()" out = mHCExpandCombineOp.apply( f, bias, @@ -196,41 +379,83 @@ def mhc_fused_expand_combine( H_res, n, use_tf32, + fused_grad_x_acc_buffer, ) return out -def mhc_fused_projection(x: torch.Tensor, phi: torch.Tensor, use_tf32: bool = True): +def mhc_fused_projection( + x: torch.Tensor, + phi: torch.Tensor, + use_tf32: bool = True, + norm_weight: Optional[torch.Tensor] = None, + fused_grad_x_acc_buffer: Optional[torch.Tensor] = None, + use_split_k: bool = False, +): """ Fused projection operation to compute H matrices and mean square for RMSNorm (see eq. 14-15, section 4.3.1 of the DeepSeek mHC paper): H = x @ phi^T: (M, K) @ (K, N) -> (M, N), which is padded to (M, 32) for better memory access pattern in the next kernels. ms = mean(x^2, dim=-1): (M,) + If norm_weight is provided, it will be absorbed into phi. In this case, the operation becomes: + Projection: + - H = x @ (phi.T * norm_weight) = x @ phi.T * norm_weight + - ms = mean(x^2, dim=-1) + - H = H / sqrt(ms) = x @ (phi.T * norm_weight) / sqrt(ms), where this step is fused into `mhc_fused_scale` + which is equivalent to performing the computation in the normal order: + - x_normalized = RMSNorm(x) = x * norm_weight / sqrt(ms) + - H = x_normalized @ phi.T = (x / sqrt(ms) @ phi.T) * norm_weight + Note: the current implementation only supports n=4 Parameters ---------- x : torch.Tensor input tensor of shape (M, K), where M=s*b is the batch size and K=nC is the hidden dimension after expansion. + dtype is torch.bfloat16 or torch.float32 phi : torch.Tensor projection matrix of shape (N, K), where N=2n+n*n (=24 for n=4) + dtype is torch.bfloat16 or torch.float32 use_tf32 : bool - whether to use TF32 precision for matmul operations. If False, it will use ieee for better precision. + whether to use TF32 precision for matmul operations. If False, it will use IEEE or TF32x3 for better precision. + Due to a triton bug (https://github.com/triton-lang/triton/issues/10176), `mhc_fused_projection` will use TF32x3 if x is bf16 and phi is fp32, or use IEEE if otherwise. This is mainly used by our unittests since TF32 precision will introduce some errors and cause tests to fail. + norm_weight : torch.Tensor or None + optional, the weight for RMSNorm, of shape (K,), which is the learnable per-element affine parameters (gamma) applied to RMSNorm + dtype is torch.bfloat16 or torch.float32 + fused_grad_x_acc_buffer : Optional[torch.Tensor] + A pre-allocated buffer for inplace gradient accumulation to avoid PyTorch autograd overhead. + If not None, triton kernels will accumulate the gradient of x into this same buffer to avoid copying the gradient by PyTorch. + This optimization requires the operation order to be mhc_fused_projection -> mhc_fused_aggregate -> mhc_fused_expand_combine. + Note: the buffer must have dtype float32, and it will be cast to the activation's dtype and be returned in mhc_fused_projection + use_split_k : bool + whether to use split-K reduction with atomic adds for the projection. Faster for large K, but + non-deterministic; requires NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 Returns ------- H : torch.Tensor - Projected matrix of shape (M, 32), where only the first N elements in the last dimension are valid. + Projected matrix of shape (M, 32), where only the first N elements in the last dimension are valid, + with dtype float32 ms : torch.Tensor - Mean square of shape (M,), which is used for RMSNorm in the next kernel. + Mean square of shape (M,), which is used for RMSNorm in the next kernel, + with dtype float32 """ assert ( phi.shape[0] == 24 ), "Currently only n=4 is supported, which means phi should have 24 in its first dimension" - check_deterministic("mhc_fused_projection") - H, ms = mHCProjectionOp.apply(x, phi, use_tf32) + check_deterministic("mhc_fused_projection", use_split_k) + if fused_grad_x_acc_buffer is not None: + assert ( + fused_grad_x_acc_buffer.dtype == torch.float32 + ), "fused_grad_x_acc_buffer must be fp32" + assert ( + fused_grad_x_acc_buffer.numel() == x.numel() + ), "fused_grad_x_acc_buffer.numel() must match x.numel()" + H, ms = mHCProjectionOp.apply( + x, phi, norm_weight, use_tf32, fused_grad_x_acc_buffer, use_split_k + ) return H, ms @@ -240,16 +465,29 @@ class mHCProjectionOp(torch.autograd.Function): """ @staticmethod - def forward(ctx, x, phi, use_tf32=True): + def forward( + ctx, + x, + phi, + norm_weight=None, + use_tf32=True, + fused_grad_x_acc_buffer=None, + use_split_k=False, + ): """ The forward pass of the fused projection operation. Computes H = x @ phi^T and the mean + If norm_weight is provided, it will be absorbd by phi square ms = mean(x^2, dim=-1) for RMSNorm in a single fused kernel. Parameters: ctx : The context object. x (tensor): The input tensor of shape (M, K), where M=s*b is the flattened batch dimension and K=nC is the hidden dimension after expansion. phi (tensor): The projection matrix of shape (N, K), where N=2n+n*n (=24 for n=4). - use_tf32 (bool): Whether to use TF32 precision for matmul operations. If False, uses IEEE for better precision. + norm_weight (tensor or None): Optional, or tensor of shape (K,). RMSNorm's learnable per-element affine parameters + use_tf32 (bool): Whether to use TF32 precision for matmul operations. If False, uses IEEE or TF32x3 for better precision. + Due to a triton bug (https://github.com/triton-lang/triton/issues/10176), we will use TF32x3 if x is bf16 and phi is fp32, or use IEEE if otherwise. + n (int): Number of hyper connections, where only n=4 is supported in the current implementation. + fused_grad_x_acc_buffer (torch.Tensor or None): A pre-allocated buffer for inplace gradient accumulation to avoid PyTorch autograd overhead. Returns: tuple: A tuple of (H, ms) where H is the projected matrix of shape (M, 32) padded for memory alignment (only the first N elements are valid), and ms is the mean square of shape (M,) in FP32. @@ -267,9 +505,7 @@ def forward(ctx, x, phi, use_tf32=True): # Pad H to (s, b, 32) for better memory access pattern in the kernel, but only the first N elements in the last dimension are valid H = torch.zeros((M, 32), device=device, dtype=torch.float32) - ms = torch.zeros( - (M,), device=device, dtype=torch.float32 - ) # Mean square for x, used to compute RMSNorm in the next kernel + ms = torch.zeros((M,), device=device, dtype=torch.float32) # pylint: disable=unnecessary-lambda-assignment grid = lambda META: ( @@ -277,6 +513,31 @@ def forward(ctx, x, phi, use_tf32=True): triton.cdiv(K, META["BLOCK_SIZE_K"]), ) + ctx.save_for_backward(x, phi, ms, norm_weight) + ctx.phi_dtype = phi.dtype + ctx.fused_grad_x_acc_buffer = fused_grad_x_acc_buffer + ctx.use_split_k = use_split_k + + if norm_weight is not None: + phi = phi * norm_weight.to(torch.float32) + elif not use_tf32 and x.dtype == torch.bfloat16 and phi.dtype == torch.bfloat16: + # tl.dot ignores input_precision when both operands are bf16 and always uses the + # bf16 MMA, Upcast phi so the dot becomes fp32 x fp32 and honors the tf32x3 precision selected below. + phi = phi.to(torch.float32) + + use_tma = _support_tma(x) and _tma_aligned(x) and _tma_aligned(phi) + if use_tma: + _init_tma_allocator() + + precision = "tf32" if ctx.use_tf32 else "ieee" + # If upcasting from bf16 to fp32 takes place inside the triton kernel, triton will ignore "ieee" precision and use tf32 anyway + # See https://github.com/triton-lang/triton/issues/10176 for detail. + # Therefore, we need to use tf32x3 instead which at least has better accuracy than tf32 just to make the tests pass. In production + # precision should be tf32 so it's not affected. + if precision == "ieee" and x.dtype == torch.bfloat16 and phi.dtype == torch.float32: + precision = "tf32x3" + ctx.precision = precision + _mhc_projection_fwd_fused[grid]( x_ptr=x, # (M, K) phi_ptr=phi, # (N, K) @@ -292,22 +553,31 @@ def forward(ctx, x, phi, use_tf32=True): stride_hm=32, stride_hn=1, stride_ms=1, + stride_norm_weight=1, BLOCK_SIZE_N=32, - precision="tf32" if use_tf32 else "ieee", + precision=precision, + USE_SPLIT_K=use_split_k, + USE_TMA=use_tma, ) - ctx.save_for_backward(x, phi, ms) - ctx.phi_dtype = phi.dtype - - return H.to(ctx.dtype), ms # Keep ms in fp32 + return H, ms # Keep both in fp32, which will be passed to sigmoid in mHCScaleFusedOp @staticmethod def backward(ctx, grad_H, grad_ms): """ The backward pass of the fused projection operation. Computes gradients for x and phi. - grad_phi = grad_H^T @ x, truncated to the first N rows. - grad_x = grad_H @ phi + 2 * x * grad_ms / K, where the second term is the gradient contribution from + - grad_psi = grad_H^T @ x: (2n + n^2, M) @ (M, nC) = (2n + n^2, nC), where grad_H's last dim is padded to 32 + If norm_weight is None: + - grad_phi = grad_psi + Otherwise, + - grad_phi = grad_psi * norm_weight: (2n + n^2, nC) * (nC,) = (2n + n^2, nC) + - grad_norm_weight = sum(grad_psi * phi, dim=0): ((2n + n^2, nC) * (2n + n^2, nC)).sum(dim=0) -> (nC,) + Reorder a bit: + - grad_phi = grad_H^T @ x * norm_weight + - grad_norm_weight = sum((grad_H^T @ x) * phi, dim=0) + + - grad_x = grad_H @ phi + 2 * x * grad_ms / K, where the second term is the gradient contribution from the mean square computation fused in the forward pass. Parameters: @@ -316,9 +586,9 @@ def backward(ctx, grad_H, grad_ms): grad_ms (tensor): The gradient of the loss with respect to the mean square, of shape (M,). Returns: - tuple: A tuple with the gradients (grad_x, grad_phi, None). + tuple: A tuple with the gradients (grad_x, grad_phi, grad_norm_weight, None). """ - x, phi, ms = ctx.saved_tensors + x, phi, ms, norm_weight = ctx.saved_tensors M, K = x.shape device = x.device @@ -332,12 +602,63 @@ def backward(ctx, grad_H, grad_ms): M, ) - grad_x = torch.empty((M, K), device=device, dtype=x.dtype) + if ctx.fused_grad_x_acc_buffer is not None: + grad_x = ctx.fused_grad_x_acc_buffer.view_as(x) + else: + grad_x = torch.empty((M, K), device=device, dtype=x.dtype) + + if norm_weight is not None: + # With norm_weight, we need a fused kernel to perform GEMM and output both phi & norm_weight gradients + # pylint: disable=unnecessary-lambda-assignment + grid = lambda META: ( + triton.cdiv(K, META["BLOCK_SIZE_K"]), + triton.cdiv(M, META["BLOCK_SIZE_M"]), + ) - grad_x = torch.empty((M, K), device=device, dtype=x.dtype) - grad_phi = general_gemm(x, grad_H, out_dtype=torch.float32, layout="NT")[0][:N, :].to( - phi.dtype - ) # (2n + n^2, M) @ (M, nC) = (2n + n^2, nC); grad_H's last dim is padded to 32 + # For reduction over M, we should prefer parallelizing over M since it's likely to be better, unless determinism is enforced + if ctx.use_split_k: + # atomic_add accumulation needs zeroed fp32 accumulators + grad_phi = torch.zeros_like(phi, dtype=torch.float32) + grad_norm_weight = torch.zeros_like(norm_weight, dtype=torch.float32) + else: + # Otherwise we don't need zeroed fp32 buffer since we don't do atomic add in this path + grad_phi = torch.empty_like(phi) + grad_norm_weight = torch.empty_like(norm_weight) + + _mhc_projection_bwd_fused_dphi[grid]( + x_ptr=x, # (M, K) + grad_H_ptr=grad_H, # (M, 32) + phi_ptr=phi, # (N, K) + norm_weight_ptr=norm_weight, # (K,) + grad_phi_ptr=grad_phi, # (N, K) + grad_norm_weight_ptr=grad_norm_weight, # (K,) + M=M, + N=N, + K=K, + stride_xm=K, + stride_xk=1, + stride_grad_Hm=32, + stride_grad_Hn=1, + stride_phin=K, + stride_phik=1, + stride_norm_weight=1, + stride_grad_phin=K, + stride_grad_phik=1, + stride_grad_norm_weight=1, + BLOCK_SIZE_N=32, + precision="tf32" if ctx.use_tf32 else "ieee", + USE_SPLIT_M=ctx.use_split_k, + ) + + grad_phi = grad_phi.to(phi.dtype) + grad_norm_weight = grad_norm_weight.to(norm_weight.dtype) + else: + # Without norm_weight, this is only a GEMM with no fusion needed so we let cuBLAS handle it + grad_phi = general_gemm( + x.to(grad_H.dtype), grad_H, out_dtype=torch.float32, layout="NT" + )[0][:N, :] + grad_phi = grad_phi.to(phi.dtype) + grad_norm_weight = None # pylint: disable=unnecessary-lambda-assignment grid = lambda META: ( @@ -345,10 +666,11 @@ def backward(ctx, grad_H, grad_ms): triton.cdiv(K, META["BLOCK_SIZE_K"]), ) - _mhc_projection_bwd_fused[grid]( + _mhc_projection_bwd_fused_dx[grid]( x_ptr=x, grad_x_ptr=grad_x, # (M, K) phi_ptr=phi, # (N, K) + norm_weight_ptr=norm_weight, # (K,) grad_h_ptr=grad_H, # (M, 32) grad_ms_ptr=grad_ms, # (M,) M=M, @@ -360,16 +682,19 @@ def backward(ctx, grad_H, grad_ms): stride_grad_xk=1, stride_phin=K, stride_phik=1, + stride_norm_weight=1, stride_grad_phin=K, stride_grad_phik=1, stride_grad_hm=32, stride_grad_hn=1, stride_grad_ms=1, BLOCK_SIZE_N=32, - precision="tf32" if ctx.use_tf32 else "ieee", + precision=ctx.precision, + FUSE_GRAD_X_ACC=ctx.fused_grad_x_acc_buffer is not None, + HAS_NORM_WEIGHT=norm_weight is not None, ) - return grad_x.to(ctx.dtype), grad_phi.to(ctx.dtype), None + return grad_x.to(x.dtype), grad_phi, grad_norm_weight, None, None, None class mHCScaleFusedOp(torch.autograd.Function): @@ -467,12 +792,24 @@ def backward(ctx, grad_out): grad_h = torch.zeros( (M, 32), device=grad_out.device, dtype=grad_out.dtype ) # Pad the grad_h to 32 in the last dimension + grad_ms = torch.zeros((M,), device=grad_out.device, dtype=grad_out.dtype) + + # grad_a and grad_b are reductions over the M dimension, which is split across grid blocks. + # scale_config fixes BLOCK_SIZE_M=128, so the grid over M is deterministic. + BLOCK_SIZE_M = 128 + grid_m = triton.cdiv(M, BLOCK_SIZE_M) + grad_alpha = torch.zeros((3,), device=grad_out.device, dtype=grad_out.dtype) grad_beta_padded = torch.zeros((1, 32), device=grad_out.device, dtype=grad_out.dtype) grad_beta = grad_beta_padded[ :, :N ] # Use only the first N elements for grad_beta, the rest are just padding - grad_ms = torch.zeros((M,), device=grad_out.device, dtype=grad_out.dtype) + if ENFORCE_DETERMINISTIC: + ws_grad_a = torch.empty((grid_m, 4), device=grad_out.device, dtype=torch.float32) + ws_grad_b = torch.empty((grid_m, 32), device=grad_out.device, dtype=torch.float32) + else: + ws_grad_a = None + ws_grad_b = None # pylint: disable=unnecessary-lambda-assignment grid = lambda META: (triton.cdiv(M, META["BLOCK_SIZE_M"]),) @@ -487,6 +824,8 @@ def backward(ctx, grad_out): grad_b_ptr=grad_beta, grad_ms_ptr=grad_ms, ms_ptr=ms, + ws_grad_a_ptr=ws_grad_a, + ws_grad_b_ptr=ws_grad_b, M=M, n=n, stride_grad_out_m=32, @@ -504,13 +843,19 @@ def backward(ctx, grad_out): stride_ms=1, BLOCK_SIZE_N=32, eps=torch.finfo(ms.dtype).eps, + DETERMINISTIC=ENFORCE_DETERMINISTIC, ) + if ENFORCE_DETERMINISTIC: + grad_alpha = ws_grad_a.sum(dim=0)[:3] # Sum partials across blocks; first 3 are grad_a + grad_beta_padded = ws_grad_b.sum(dim=0, keepdim=True) # Sum partials across blocks + grad_beta = grad_beta_padded[:, :N] + return ( - grad_h.to(ctx.dtype), - grad_alpha.to(ctx.dtype), - grad_beta.to(ctx.dtype), - grad_ms.to(ctx.dtype), + grad_h, + grad_alpha.to(alpha.dtype), + grad_beta.to(alpha.dtype), # We assume alpha and beta have the same dtype + grad_ms, None, ) @@ -676,7 +1021,6 @@ def backward(ctx, grad_out): ) grad_res = grad_res.view(s, b, n, n) - return grad_res.to(ctx.dtype), None, None, None @@ -686,7 +1030,7 @@ class mHCAggregateOp(torch.autograd.Function): """ @staticmethod - def forward(ctx, x, H_pre, n, use_tf32=True): + def forward(ctx, x, H_pre, n, use_tf32=True, fused_grad_x_acc_buffer=None): """ The forward pass of the aggregate operation. Merges n activation streams into one by computing a weighted sum using H_pre: @@ -699,6 +1043,7 @@ def forward(ctx, x, H_pre, n, use_tf32=True): H_pre (tensor): The pre-connection matrix of shape (s, b, n), used as weights for aggregation. n (int): The number of hyper connections (only n=4 is supported). use_tf32 (bool): Whether to use TF32 precision for matmul operations. + fused_grad_x_acc_buffer (torch.Tensor or None): A pre-allocated buffer for inplace gradient accumulation to avoid PyTorch autograd overhead. Returns: tensor: The aggregated output of shape (s, b, C). @@ -735,6 +1080,7 @@ def forward(ctx, x, H_pre, n, use_tf32=True): ctx.save_for_backward(x, H_pre) ctx.n = n ctx.use_tf32 = use_tf32 + ctx.fused_grad_x_acc_buffer = fused_grad_x_acc_buffer return out @@ -763,7 +1109,11 @@ def backward(ctx, grad_output): assert n == 4, "Only n=4 is supported in this implementation" M = s * b - grad_x = torch.empty_like(x) + if ctx.fused_grad_x_acc_buffer is not None: + grad_x = ctx.fused_grad_x_acc_buffer.view_as(x) + else: + grad_x = torch.empty_like(x) + grad_H_pre = torch.zeros( (s, b, n), dtype=torch.float32, device=H_pre.device ) # We need to use atomic_add for this so we need higher precision @@ -790,11 +1140,15 @@ def backward(ctx, grad_output): stride_grad_xm=nC, stride_grad_xCn=1, precision="tf32" if ctx.use_tf32 else "ieee", + FUSE_GRAD_X_ACC=ctx.fused_grad_x_acc_buffer is not None, ) grad_H_pre = grad_H_pre.to(H_pre.dtype) # Cast back to the original dtype of H_pre - return grad_x, grad_H_pre, None, None + if ctx.fused_grad_x_acc_buffer is not None: + grad_x = None + + return grad_x, grad_H_pre, None, None, None class mHCExpandCombineOp(torch.autograd.Function): @@ -803,7 +1157,7 @@ class mHCExpandCombineOp(torch.autograd.Function): """ @staticmethod - def forward(ctx, f, bias, H_post, x, H_res, n, use_tf32=True): + def forward(ctx, f, bias, H_post, x, H_res, n, use_tf32=True, fused_grad_x_acc_buffer=None): """ The forward pass of the expand and combine operation. Expands the sub-layer output f back to n streams using H_post, and combines with the residual connections using H_res: @@ -819,6 +1173,7 @@ def forward(ctx, f, bias, H_post, x, H_res, n, use_tf32=True): H_res (tensor): The residual connection matrix of shape (s, b, n, n). n (int): The number of hyper connections (only n=4 is supported). use_tf32 (bool): Whether to use TF32 precision for matmul operations. + fused_grad_x_acc_buffer (torch.Tensor or None): A pre-allocated buffer for inplace gradient accumulation to avoid PyTorch autograd overhead. Returns: tensor: The expanded and combined output of shape (s, b, C, n). @@ -843,45 +1198,29 @@ def forward(ctx, f, bias, H_post, x, H_res, n, use_tf32=True): triton.cdiv(M, META["BLOCK_SIZE_M"]), ) - if bias is None: - _mhc_expand_combine_fwd[grid]( - f_ptr=f, - H_post_ptr=H_post, - x_ptr=x, - H_res_ptr=H_res, - output_ptr=out, - M=M, - C=C, - n=n, - stride_fm=C, - stride_fc=1, - stride_xm=Cn, - stride_xCn=1, - stride_output_m=Cn, - stride_output_Cn=1, - ) - else: - _mhc_expand_combine_with_bias_fwd[grid]( - f_ptr=f, - bias_ptr=bias, - H_post_ptr=H_post, - x_ptr=x, - H_res_ptr=H_res, - output_ptr=out, - M=M, - C=C, - n=n, - stride_fm=C, - stride_fc=1, - stride_bias=1, - stride_xm=Cn, - stride_xCn=1, - stride_output_m=Cn, - stride_output_Cn=1, - ) + _mhc_expand_combine_fwd[grid]( + f_ptr=f, + bias_ptr=bias, + H_post_ptr=H_post, + x_ptr=x, + H_res_ptr=H_res, + output_ptr=out, + M=M, + C=C, + n=n, + stride_fm=C, + stride_fc=1, + stride_bias=1, + stride_xm=Cn, + stride_xCn=1, + stride_output_m=Cn, + stride_output_Cn=1, + HAS_BIAS=bias is not None, + ) ctx.n = n ctx.have_bias = bias is not None + ctx.fused_grad_x_acc_buffer = fused_grad_x_acc_buffer if bias is not None: ctx.save_for_backward(f, bias, H_post, x, H_res) else: @@ -919,14 +1258,26 @@ def backward(ctx, grad_output): M = s * b grad_f = torch.empty_like(f) - grad_bias = torch.zeros_like(bias, dtype=torch.float32) if bias is not None else None - grad_H_post = torch.zeros_like( - H_post, dtype=torch.float32 - ) # We need to use atomic_add for this so we need higher precision - grad_x = torch.empty_like(x) - grad_H_res = torch.zeros_like( - H_res, dtype=torch.float32 - ) # We need to use atomic_add for this so we need higher precision + if ctx.fused_grad_x_acc_buffer is not None: + grad_x = ctx.fused_grad_x_acc_buffer.view_as(x) + else: + grad_x = torch.empty_like(x) + + # Since triton's autotune will reset grad_bias pointer when tuning, we need an empty placeholder here + grad_bias = torch.empty(1, device=grad_output.device, dtype=grad_output.dtype) + grad_H_post = torch.empty_like(H_post) + grad_H_res = torch.empty_like(H_res) + + # grad_bias is a reduction over M. In deterministic mode each block writes its partial to a + # workspace row (reduced in torch); otherwise it is atomic-added into grad_bias directly. + # The deterministic (non-split) path fixes BLOCK_SIZE_M=4, so the grid over M is cdiv(M, 4). + grad_bias_ws = None + if bias is not None: + grad_bias = torch.zeros_like(bias, dtype=torch.float32) + if ENFORCE_DETERMINISTIC: + grad_bias_ws = torch.empty( + (triton.cdiv(M, 4), C), device=grad_output.device, dtype=torch.float32 + ) # pylint: disable=unnecessary-lambda-assignment grid = lambda META: ( @@ -934,66 +1285,53 @@ def backward(ctx, grad_output): triton.cdiv(M, META["BLOCK_SIZE_M"]), ) - if bias is None: - _mhc_expand_combine_bwd[grid]( - grad_output_ptr=grad_output, - f_ptr=f, - H_post_ptr=H_post, - x_ptr=x, - H_res_ptr=H_res, - grad_H_post_ptr=grad_H_post, - grad_f_ptr=grad_f, - grad_H_res_ptr=grad_H_res, - grad_x_ptr=grad_x, - M=M, - C=C, - n=n, - stride_grad_output_m=n * C, - stride_grad_output_Cn=1, - stride_fm=C, - stride_fc=1, - stride_xm=n * C, - stride_xCn=1, - stride_grad_fm=C, - stride_grad_fc=1, - stride_grad_xm=n * C, - stride_grad_xCn=1, - precision="tf32" if ctx.use_tf32 else "ieee", - ) - else: - _mhc_expand_combine_with_bias_bwd[grid]( - grad_output_ptr=grad_output, - f_ptr=f, - bias_ptr=bias, - H_post_ptr=H_post, - x_ptr=x, - H_res_ptr=H_res, - grad_H_post_ptr=grad_H_post, - grad_f_ptr=grad_f, - grad_bias_ptr=grad_bias, - grad_H_res_ptr=grad_H_res, - grad_x_ptr=grad_x, - M=M, - C=C, - n=n, - stride_grad_output_m=n * C, - stride_grad_output_Cn=1, - stride_fm=C, - stride_fc=1, - stride_bias=1, - stride_xm=n * C, - stride_xCn=1, - stride_grad_fm=C, - stride_grad_fc=1, - stride_grad_bias=1, - stride_grad_xm=n * C, - stride_grad_xCn=1, - precision="tf32" if ctx.use_tf32 else "ieee", - ) + _mhc_expand_combine_bwd[grid]( + grad_output_ptr=grad_output, + f_ptr=f, + bias_ptr=bias, + H_post_ptr=H_post, + x_ptr=x, + H_res_ptr=H_res, + grad_H_post_ptr=grad_H_post, + grad_f_ptr=grad_f, + grad_bias_ptr=grad_bias, + grad_bias_ws_ptr=grad_bias_ws, + grad_H_res_ptr=grad_H_res, + grad_x_ptr=grad_x, + M=M, + C=C, + n=n, + stride_grad_output_m=n * C, + stride_grad_output_Cn=1, + stride_fm=C, + stride_fc=1, + stride_bias=1, + stride_xm=n * C, + stride_xCn=1, + stride_grad_fm=C, + stride_grad_fc=1, + stride_grad_bias=1, + stride_grad_bias_ws_m=C, + stride_grad_bias_ws_c=1, + stride_grad_xm=n * C, + stride_grad_xCn=1, + precision="tf32" if ctx.use_tf32 else "ieee", + HAS_BIAS=bias is not None, + FUSE_GRAD_X_ACC=ctx.fused_grad_x_acc_buffer is not None, + DETERMINISTIC=ENFORCE_DETERMINISTIC, + ) grad_H_post = grad_H_post.to(H_post.dtype) # Cast back to the original dtype of H_post grad_H_res = grad_H_res.to(H_res.dtype) # Cast back to the original dtype of H_res - if bias is not None: + if bias is None: + # If no bias, replace the grad_bias placeholder with None + grad_bias = None + elif ENFORCE_DETERMINISTIC: + grad_bias = grad_bias_ws.sum(dim=0).to(bias.dtype) # Sum partials across blocks + else: grad_bias = grad_bias.to(bias.dtype) - return grad_f, grad_bias, grad_H_post, grad_x, grad_H_res, None, None + if ctx.fused_grad_x_acc_buffer is not None: + grad_x = None + + return grad_f, grad_bias, grad_H_post, grad_x, grad_H_res, None, None, None diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index 4cd9583aa5..4fb6ff6a4f 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -88,6 +88,7 @@ def _get_device_compute_capability(device: torch.device) -> Tuple[int, int]: return (props.major, props.minor) +@torch.compiler.assume_constant_result def get_device_compute_capability() -> Tuple[int, int]: """CUDA compute capability of current GPU""" return _get_device_compute_capability(torch.cuda.current_device()) @@ -693,7 +694,7 @@ def is_non_tn_fp8_gemm_supported() -> bool: @functools.lru_cache(maxsize=None) -def get_cudnn_version() -> Tuple[int, int, int]: +def _get_cudnn_version() -> Tuple[int, int, int]: """Runtime cuDNN version (major, minor, patch)""" # ROCm fused attn does not use cudnn, return high numbers to avoid tests filtering out if IS_HIP_EXTENSION: @@ -707,6 +708,12 @@ def get_cudnn_version() -> Tuple[int, int, int]: return (major, minor, patch) +@torch.compiler.assume_constant_result +def get_cudnn_version() -> Tuple[int, int, int]: + """Runtime cuDNN version (major, minor, patch)""" + return _get_cudnn_version() + + def canonicalize_device(device: Optional[torch.device | str]) -> torch.device: """Canonicalize PyTorch device