diff --git a/build_tools/VERSION.txt b/build_tools/VERSION.txt index 5e6b220a346..38018b6482d 100644 --- a/build_tools/VERSION.txt +++ b/build_tools/VERSION.txt @@ -1 +1 @@ -2.20.0.dev0 +2.21.0-dev0 diff --git a/build_tools/jax.py b/build_tools/jax.py index 031432e6f90..10906dc708a 100644 --- a/build_tools/jax.py +++ b/build_tools/jax.py @@ -5,6 +5,7 @@ """JAX related extensions.""" import os +import warnings from pathlib import Path from packaging import version @@ -95,16 +96,24 @@ def setup_jax_extension( if (discovered_nccl_include_path := nccl_include_path()) is not None: include_dirs.append(discovered_nccl_include_path) include_dirs.append(cudnn_frontend_include_path()) + xla_include_path = xla_path() include_dirs.extend( [ common_header_files, common_header_files / "common", common_header_files / "common" / "include", csrc_header_files, - xla_path(), + xla_include_path, ] ) + # Match the borrowed-comm path's compile-time header check. + if not (Path(xla_include_path) / "xla/ffi/api/collectives_c_api.h").is_file(): + warnings.warn( + f"XLA headers in {xla_include_path} do not include " + "xla/ffi/api/collectives_c_api.h; the EP borrowed-comm path will not be built." + ) + # Compile flags cxx_flags = ["-O3"] if debug_build_enabled(): diff --git a/qa/L2_jax_distributed_unittest/test.sh b/qa/L2_jax_distributed_unittest/test.sh index 330b254e7db..c64a1e561d9 100644 --- a/qa/L2_jax_distributed_unittest/test.sh +++ b/qa/L2_jax_distributed_unittest/test.sh @@ -15,4 +15,6 @@ mkdir -p "$XML_LOG_DIR" 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_* # NCCL EP multi-process suite. The launcher skips when fewer than 4 GPUs or no NVLink is detected. +# Runs the borrowed-comm suite too (L2 only). +export NVTE_JAX_UNITTEST_LEVEL="L2" TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/multi_process_launch_ep.sh diff --git a/tests/jax/multi_process_launch_ep.sh b/tests/jax/multi_process_launch_ep.sh index 8547d77f2b4..eb1735fef96 100755 --- a/tests/jax/multi_process_launch_ep.sh +++ b/tests/jax/multi_process_launch_ep.sh @@ -7,6 +7,16 @@ SCRIPT_NAMES="${SCRIPT_NAMES:-test_multi_process_ep.py}" TEST_TIMEOUT_S="${TEST_TIMEOUT_S:-180}" +# Each communicator mode needs a fresh process group. +if [ -z "${NVTE_TEST_EP_CLASSES:-}" ]; then + RET=0 + NVTE_TEST_EP_CLASSES="TestEP,TestEPOverflowDrop,TestEpDomainGrouping" \ + bash "${BASH_SOURCE[0]}" || RET=1 + NVTE_TEST_EP_CLASSES="TestEPBorrowedComm" \ + bash "${BASH_SOURCE[0]}" || RET=1 + exit "$RET" +fi + XLA_BASE_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true --xla_gpu_graph_min_graph_size=1" diff --git a/tests/jax/test_multi_process_ep.py b/tests/jax/test_multi_process_ep.py index 47af0b0c397..e26d3e01a9d 100644 --- a/tests/jax/test_multi_process_ep.py +++ b/tests/jax/test_multi_process_ep.py @@ -25,6 +25,7 @@ import re import sys import unittest +from unittest import mock import jax import jax.experimental.multihost_utils as jmu @@ -47,8 +48,13 @@ ep_dispatch_fwd, ep_combine_fwd, get_ep_config, + is_ep_borrowed_comm_built, + use_nccl_comm_from_xla, +) +from transformer_engine.jax.version_utils import ( + is_collective_stream_supported, + is_xla_ffi_collectives_supported, ) -from transformer_engine.jax.version_utils import is_collective_stream_supported # ── Test config ───────────────────────────────────────────────────────────── @@ -107,11 +113,23 @@ def _local_device_sm(): class TestEP(unittest.TestCase): + # Selects the EP comm path for this class. False forces the self-hosted NCCL + # comm; the TestEPBorrowedComm subclass flips it to exercise the borrowed path. + USE_BORROWED_COMM = False + @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})") + if cls.USE_BORROWED_COMM and not ( + is_ep_borrowed_comm_built() and is_xla_ffi_collectives_supported() + ): + raise unittest.SkipTest("EP borrowed-comm path needs a newer JAX/XLA build") + cls._prev_comm_env = os.environ.get("NVTE_JAX_EP_NCCL_COMM_FROM_XLA") + os.environ["NVTE_JAX_EP_NCCL_COMM_FROM_XLA"] = "1" if cls.USE_BORROWED_COMM else "0" + # Drop any communicator a prior class left so we bootstrap on a clean slate. + ep_finalize() cls.num_procs = jax.process_count() cls.rank = jax.process_index() cls.dp, cls.ep = _factor_dp_ep(cls.num_procs) @@ -144,6 +162,15 @@ def setUpClass(cls): # alignment exercises dispatch_output_per_expert_alignment end-to-end. cls.hk = EpLayerConfig(top_k=TOP_K, dispatch_output_per_expert_alignment=16) + @classmethod + def tearDownClass(cls): + # Leave a clean slate for the next class and restore the env override. + ep_finalize() + if cls._prev_comm_env is None: + os.environ.pop("NVTE_JAX_EP_NCCL_COMM_FROM_XLA", None) + else: + os.environ["NVTE_JAX_EP_NCCL_COMM_FROM_XLA"] = cls._prev_comm_env + # ── Bootstrap precondition ──────────────────────────────────────────── def test_bootstrap_rejects_missing_ep_axis(self): @@ -820,6 +847,37 @@ def bwd_only(eo, toks, idx, w, g): self.assertEqual(hlo.count(op), 0, f"unexpected XLA {op} in bwd HLO:\n{hlo}") +# ── Borrowed-comm path ─────────────────────────────────────────────────────── + + +class TestEPBorrowedComm(TestEP): + """Re-run EP primitives on the XLA borrowed-comm path. + + Skipped entirely unless the build and installed JAX both provide the + collectives FFI extension. To keep L0/L1 fast, only a small smoke subset + (_SMOKE) runs by default; the full borrowed-path suite runs at L2 + (NVTE_JAX_UNITTEST_LEVEL=L2). + """ + + USE_BORROWED_COMM = True + + # Representative cases kept outside L2: one dispatch/combine round-trip (fwd) + # and its gradient (bwd). Every other inherited case runs only at L2. + _SMOKE = frozenset( + { + "test_primitive_dispatch_combine_identity_uniform", + "test_primitive_dispatch_combine_identity_bwd_uniform", + } + ) + + def setUp(self): + if ( + os.environ.get("NVTE_JAX_UNITTEST_LEVEL", "L0") != "L2" + and self._testMethodName not in self._SMOKE + ): + self.skipTest("borrowed-comm full suite runs at L2 (NVTE_JAX_UNITTEST_LEVEL=L2)") + + # ── Drop-on-overflow ───────────────────────────────────────────────────────── @@ -846,6 +904,8 @@ 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._prev_comm_env = os.environ.get("NVTE_JAX_EP_NCCL_COMM_FROM_XLA") + os.environ["NVTE_JAX_EP_NCCL_COMM_FROM_XLA"] = "0" cls.num_procs = jax.process_count() cls.rank = jax.process_index() cls.dp, cls.ep = _factor_dp_ep(cls.num_procs) @@ -873,6 +933,10 @@ def setUpClass(cls): def tearDownClass(cls): # Leave a clean slate so another class can bootstrap after us. ep_finalize() + if cls._prev_comm_env is None: + os.environ.pop("NVTE_JAX_EP_NCCL_COMM_FROM_XLA", None) + else: + os.environ["NVTE_JAX_EP_NCCL_COMM_FROM_XLA"] = cls._prev_comm_env def _make_concentrated_inputs(self): """All top-1 routes to expert 0; top-2 spread over the rest, so the rank @@ -961,6 +1025,65 @@ def test_ep_tp_splits_domains(self): self.assertEqual(domains, {0: [0, 2, 4, 6], 1: [1, 3, 5, 7]}) +# ── Comm-path selection (single-process; no GPU needed) ────────────────────── + + +class TestEpCommSelection(unittest.TestCase): + """use_nccl_comm_from_xla() build/version gating and NVTE_JAX_EP_NCCL_COMM_FROM_XLA override.""" + + @staticmethod + def _use(env, built, supported): + import transformer_engine.jax.cpp_extensions.ep as ep_mod + + prev = os.environ.pop("NVTE_JAX_EP_NCCL_COMM_FROM_XLA", None) + if env is not None: + os.environ["NVTE_JAX_EP_NCCL_COMM_FROM_XLA"] = env + try: + with mock.patch.object( + ep_mod, "is_ep_borrowed_comm_built", return_value=built + ), mock.patch.object( + ep_mod, "is_xla_ffi_collectives_supported", return_value=supported + ): + return ep_mod.use_nccl_comm_from_xla() + finally: + os.environ.pop("NVTE_JAX_EP_NCCL_COMM_FROM_XLA", None) + if prev is not None: + os.environ["NVTE_JAX_EP_NCCL_COMM_FROM_XLA"] = prev + + def test_auto_requires_build_and_version(self): + # Env unset: borrowed path only when both build and JAX support it. + self.assertTrue(self._use(None, built=True, supported=True)) + self.assertFalse(self._use(None, built=True, supported=False)) + self.assertFalse(self._use(None, built=False, supported=True)) + + def test_env_override_wins_over_version(self): + self.assertTrue(self._use("1", built=True, supported=False)) + self.assertFalse(self._use("0", built=True, supported=True)) + + def test_force_on_without_build_raises(self): + with self.assertRaisesRegex(RuntimeError, "without the EP borrowed-comm path"): + self._use("1", built=False, supported=True) + + +def _ep_test_cases(): + """Select test classes for one communicator mode.""" + all_test_cases = { + c.__name__: c + for c in (TestEP, TestEPBorrowedComm, TestEPOverflowDrop, TestEpDomainGrouping) + } + names = os.environ.get("NVTE_TEST_EP_CLASSES") + test_cases = ( + tuple(all_test_cases[name.strip()] for name in names.split(",")) + if names + else (TestEP, TestEPOverflowDrop, TestEpDomainGrouping) + ) + if TestEPBorrowedComm in test_cases and any( + c in test_cases for c in (TestEP, TestEPOverflowDrop) + ): + raise ValueError("Run borrowed-comm and self-hosted EP tests in separate processes.") + return test_cases + + # ── Entry point ────────────────────────────────────────────────────────────── @@ -972,6 +1095,20 @@ def test_ep_tp_splits_domains(self): coord_addr = sys.argv[1] proc_id = int(sys.argv[2]) num_procs = int(sys.argv[3]) + test_cases = _ep_test_cases() + + target = os.environ.get("TARGET_TEST") + if target: + name = target.split(".")[-1] + if not any( + hasattr(c, name) + for c in (TestEP, TestEPBorrowedComm, TestEPOverflowDrop, TestEpDomainGrouping) + ): + raise ValueError(f"Unknown EP test: {target}") + test_cases = tuple(c for c in test_cases if hasattr(c, name)) + if not test_cases: + unittest.TextTestRunner(verbosity=2).run(unittest.TestSuite()) + sys.exit(0) jax.distributed.initialize( coordinator_address=coord_addr, @@ -981,12 +1118,8 @@ def test_ep_tp_splits_domains(self): ) loader = unittest.TestLoader() - test_cases = (TestEP, TestEPOverflowDrop, TestEpDomainGrouping) - target = os.environ.get("TARGET_TEST") if target: - name = target.split(".")[-1] - cls = next((c for c in test_cases if hasattr(c, name)), TestEP) - suite = loader.loadTestsFromName(name, cls) + suite = unittest.TestSuite(loader.loadTestsFromName(name, c) for c in test_cases) else: suite = unittest.TestSuite(loader.loadTestsFromTestCase(c) for c in test_cases) runner = unittest.TextTestRunner(verbosity=2) diff --git a/tests/pytorch/distributed/test_comm_gemm_overlap.py b/tests/pytorch/distributed/test_comm_gemm_overlap.py index b00dbaccd32..aa4a2dfdfc3 100644 --- a/tests/pytorch/distributed/test_comm_gemm_overlap.py +++ b/tests/pytorch/distributed/test_comm_gemm_overlap.py @@ -33,6 +33,15 @@ # to avoid numerical tolerance issues of doing comm gemm overlap, limit the number of GPUs used MAX_GPUS_TO_USE = 4 +COMM_GEMM_QUANTIZATION_PARAMS = [ + pytest.param(False, "none", id="ub-bf16"), + pytest.param(False, "fp8", id="ub-fp8"), + pytest.param(False, "mxfp8", id="ub-mxfp8"), + pytest.param(True, "none", id="cublasmp-bf16"), + pytest.param(True, "fp8", id="cublasmp-fp8"), + pytest.param(True, "mxfp8", id="cublasmp-mxfp8"), +] + TEST_ROOT = Path(__file__).parent.resolve() NUM_PROCS: int = min(torch.cuda.device_count(), MAX_GPUS_TO_USE) LAUNCH_CMD = ["torchrun", f"--nproc_per_node={NUM_PROCS}"] @@ -102,10 +111,6 @@ def _run_gemm_with_overlap( if use_cublasmp: if not tex.nvte_built_with_cublasmp(): pytest.skip("Transformer Engine not built with cuBLASMp (NVTE_WITH_CUBLASMP=0).") - if quantization == "mxfp8": - pytest.skip( - "cuBLASMp comm+GEMM overlap does not yet support MXFP8 (block scaling)." - ) if comm_type == "RS" and not p2p and not tex.device_supports_multicast(): pytest.skip( "cuBLASMp non-P2P reduce-scatter requires NVSwitch (multicast support)." @@ -159,8 +164,6 @@ def _run_layer_with_overlap( if use_cublasmp: if not tex.nvte_built_with_cublasmp(): pytest.skip("Transformer Engine not built with cuBLASMp (NVTE_WITH_CUBLASMP=0).") - if fp8 and quantization == "mxfp8": - pytest.skip("cuBLASMp comm+GEMM overlap does not yet support MXFP8 (block scaling).") test_cmd.append("--use-cublasmp") test_env = os.environ.copy() @@ -182,8 +185,7 @@ def _run_layer_with_overlap( _assert_subprocess_succeeded(result) -@pytest.mark.parametrize("use_cublasmp", (False, True)) -@pytest.mark.parametrize("quantization", ("none", "fp8", "mxfp8")) +@pytest.mark.parametrize("use_cublasmp,quantization", COMM_GEMM_QUANTIZATION_PARAMS) @pytest.mark.parametrize("aggregate", (False, True)) def test_split_all_gather_overlaps(quantization, aggregate, use_cublasmp): """ @@ -193,8 +195,7 @@ def test_split_all_gather_overlaps(quantization, aggregate, use_cublasmp): _run_gemm_with_overlap("AG", False, True, False, aggregate, quantization, use_cublasmp) -@pytest.mark.parametrize("use_cublasmp", (False, True)) -@pytest.mark.parametrize("quantization", ("none", "fp8", "mxfp8")) +@pytest.mark.parametrize("use_cublasmp,quantization", COMM_GEMM_QUANTIZATION_PARAMS) @pytest.mark.parametrize("p2p", (False, True)) def test_split_reduce_scatter_overlaps(quantization, p2p, use_cublasmp): """ diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index 18efb6a190b..bd89ecebc9e 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -1,18 +1,21 @@ # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. -import torch +from copy import deepcopy from typing import Optional + +import pytest +import torch + from transformer_engine.pytorch.router import ( QBHistogramMode, RoutingMapFormat, fused_topk_with_score_function, fused_compute_score_for_moe_aux_loss, fused_moe_aux_loss, + mark_qb_bin_bounds_validated, ) import transformer_engine_torch as tex -import pytest -from copy import deepcopy seed = 42 torch.manual_seed(seed) @@ -621,7 +624,8 @@ def test_qb_topk_rejects_invalid_bin_bounds(histogram_mode, invalid_bounds): ) -def test_qb_topk_revalidates_updated_bin_bounds(): +@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) +def test_qb_topk_revalidates_updated_bin_bounds(histogram_mode): logits = torch.randn(8, 16, device="cuda", dtype=torch.float32) expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32) histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32) @@ -637,7 +641,7 @@ def test_qb_topk_revalidates_updated_bin_bounds(): expert_bias, qb_histogram=histogram, qb_bin_bounds=bin_bounds, - qb_histogram_mode="fused_atomic", + qb_histogram_mode=histogram_mode, ) bin_bounds.fill_(0.0) with pytest.raises(ValueError, match="finite with lower < upper"): @@ -652,7 +656,7 @@ def test_qb_topk_revalidates_updated_bin_bounds(): expert_bias, qb_histogram=histogram, qb_bin_bounds=bin_bounds, - qb_histogram_mode="fused_atomic", + qb_histogram_mode=histogram_mode, ) @@ -699,11 +703,117 @@ def test_qb_raw_binding_rejects_invalid_bin_bounds_recoverably(histogram_mode, u @pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) -def test_qb_topk_cuda_graph_uses_prevalidated_bounds(histogram_mode): +@pytest.mark.parametrize("use_dense_indices", [False, True]) +def test_qb_topk_cuda_graph_uses_mutable_bounds(histogram_mode, use_dense_indices): + logits = torch.randn(8, 16, device="cuda", dtype=torch.float32) + expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32) + histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32) + bin_bounds = torch.tensor([-1.0, 1.0], device="cuda", dtype=torch.float32) + topk_indices = ( + torch.empty(8, 4, device="cuda", dtype=torch.int32) if use_dense_indices else None + ) + + def run_router(): + return fused_topk_with_score_function( + logits, + 4, + False, + None, + None, + None, + "sigmoid", + expert_bias, + topk_indices=topk_indices, + qb_histogram=histogram, + qb_bin_bounds=bin_bounds, + qb_histogram_mode=histogram_mode, + ) + + run_router() + bounds_data_ptr = bin_bounds.data_ptr() + initial_bounds = bin_bounds.clone() + + histogram.zero_() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + probs, routing_output = run_router() + + bin_bounds.copy_(torch.tensor([-0.75, 0.25], device="cuda")) + assert bin_bounds.data_ptr() == bounds_data_ptr + histogram.zero_() + graph.replay() + torch.cuda.synchronize() + reference = qb_topk_score_function_pytorch( + logits, 4, expert_bias, bin_bounds, histogram.shape[1] + ) + torch.testing.assert_close(probs, reference["probs"]) + if use_dense_indices: + torch.testing.assert_close( + topk_indices_to_routing_map(routing_output, logits.shape[1]), + reference["routing_map"], + ) + else: + torch.testing.assert_close(routing_output, reference["routing_map"]) + torch.testing.assert_close(histogram, reference["histogram"]) + initial_reference = qb_topk_score_function_pytorch( + logits, 4, expert_bias, initial_bounds, histogram.shape[1] + ) + assert not torch.equal(histogram, initial_reference["histogram"]) + + +@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) +def test_qb_topk_cuda_graph_captures_bounds_update(histogram_mode): + logits = torch.randn(8, 16, device="cuda", dtype=torch.float32) + expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32) + histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32) + bin_bounds = torch.tensor([-1.0, 1.0], device="cuda", dtype=torch.float32) + next_bounds = torch.tensor([-0.75, 0.25], device="cuda", dtype=torch.float32) + + def run_iteration(): + probs, routing_map = fused_topk_with_score_function( + logits, + 4, + False, + None, + None, + None, + "sigmoid", + expert_bias, + qb_histogram=histogram, + qb_bin_bounds=bin_bounds, + qb_histogram_mode=histogram_mode, + ) + bin_bounds.copy_(next_bounds) + mark_qb_bin_bounds_validated(bin_bounds) + return probs, routing_map + + # Match full-iteration capture: an eager warmup ends with a bounds update, and the same + # update is part of the captured iteration. + run_iteration() + histogram.zero_() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + probs, routing_map = run_iteration() + histogram.zero_() + graph.replay() + torch.cuda.synchronize() + + reference = qb_topk_score_function_pytorch( + logits, 4, expert_bias, next_bounds, histogram.shape[1] + ) + torch.testing.assert_close(probs, reference["probs"]) + torch.testing.assert_close(routing_map, reference["routing_map"]) + torch.testing.assert_close(histogram, reference["histogram"]) + torch.testing.assert_close(bin_bounds, next_bounds) + + +@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) +def test_qb_topk_eager_call_revalidates_after_captured_bounds_update(histogram_mode): logits = torch.randn(8, 16, device="cuda", dtype=torch.float32) expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32) histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32) bin_bounds = torch.tensor([-1.0, 1.0], device="cuda", dtype=torch.float32) + next_bounds = torch.tensor([-0.75, 0.25], device="cuda", dtype=torch.float32) def run_router(): return fused_topk_with_score_function( @@ -723,11 +833,71 @@ def run_router(): run_router() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - probs, routing_map = run_router() + bin_bounds.copy_(next_bounds) + mark_qb_bin_bounds_validated(bin_bounds) + + # Replayed device writes do not advance bin_bounds._version. The capture-time marker must not + # let a subsequent eager router call inherit trust in the replay-produced values. + next_bounds.zero_() graph.replay() torch.cuda.synchronize() - assert torch.isfinite(probs).all() - assert routing_map.sum().item() == logits.shape[0] * 4 + with pytest.raises(ValueError, match="finite with lower < upper"): + run_router() + + +@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) +def test_qb_topk_cuda_graph_rejects_unvalidated_bounds(histogram_mode): + logits = torch.randn(8, 16, device="cuda", dtype=torch.float32) + expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32) + histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32) + bin_bounds = torch.tensor([-1.0, 1.0], device="cuda", dtype=torch.float32) + + with pytest.raises(RuntimeError, match="current version must be validated"): + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fused_topk_with_score_function( + logits, + 4, + False, + None, + None, + None, + "sigmoid", + expert_bias, + qb_histogram=histogram, + qb_bin_bounds=bin_bounds, + qb_histogram_mode=histogram_mode, + ) + + +@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"]) +def test_qb_topk_cuda_graph_rejects_stale_bounds_validation(histogram_mode): + logits = torch.randn(8, 16, device="cuda", dtype=torch.float32) + expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32) + histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32) + bin_bounds = torch.tensor([-1.0, 1.0], device="cuda", dtype=torch.float32) + + def run_router(): + return fused_topk_with_score_function( + logits, + 4, + False, + None, + None, + None, + "sigmoid", + expert_bias, + qb_histogram=histogram, + qb_bin_bounds=bin_bounds, + qb_histogram_mode=histogram_mode, + ) + + run_router() + bin_bounds.zero_() + with pytest.raises(RuntimeError, match="current version must be validated"): + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run_router() @pytest.mark.parametrize( diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 4436af1955d..e71ed942d4a 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -427,12 +427,12 @@ if (NVTE_WITH_CUBLASMP) NAMES nccl libnccl PATH_SUFFIXES lib REQUIRED) - # cuBLASMp 0.8 is the first release with CUDA-graph-safe overlap algos, + # cuBLASMp 0.8.1 adds MXFP8 support to the CUDA-graph-safe overlap algos, # and NCCL 2.30 is the first release with graph-safe one-sided RMA # primitives (ncclPutSignal/ncclWaitSignal) that those algos use. - if (CUBLASMP_VERSION VERSION_LESS 0.8.0) + if (CUBLASMP_VERSION VERSION_LESS 0.8.1) message(FATAL_ERROR - "NVTE_WITH_CUBLASMP requires cuBLASMp >= 0.8.0, but found cuBLASMp " + "NVTE_WITH_CUBLASMP requires cuBLASMp >= 0.8.1, but found cuBLASMp " "${CUBLASMP_VERSION} in ${CUBLASMP_INCLUDE_DIR}/cublasmp.h") endif() if (NCCL_VERSION VERSION_LESS 2.30.0) diff --git a/transformer_engine/common/comm_gemm/comm_gemm.cpp b/transformer_engine/common/comm_gemm/comm_gemm.cpp index 3545347da47..e7276496daa 100644 --- a/transformer_engine/common/comm_gemm/comm_gemm.cpp +++ b/transformer_engine/common/comm_gemm/comm_gemm.cpp @@ -260,28 +260,63 @@ void cublasmp_gemm(InitMatricesFn init_matrices_fn, NVTECommGemmCtx* ctx, NVTECo const Tensor* d, const Tensor* bias, const Tensor* pre_act_out, bool transa, bool transb, bool grad, bool accumulate, int comm_sm_count, cudaStream_t main_stream) { - for (auto t : {a, b, d}) { - NVTE_CHECK(is_tensor_scaling(t->scaling_mode), - "Unsupported scaling mode: " + std::to_string(t->scaling_mode)); + const bool tensor_scaling = + is_tensor_scaling(a->scaling_mode) && is_tensor_scaling(b->scaling_mode); + const bool mxfp8 = is_mxfp_scaling(a->scaling_mode) && is_mxfp_scaling(b->scaling_mode); + + NVTE_CHECK(tensor_scaling || mxfp8, "Unsupported scaling modes: A=", to_string(a->scaling_mode), + ", B=", to_string(b->scaling_mode)); + NVTE_CHECK(is_tensor_scaling(d->scaling_mode), + "Unsupported scaling mode for D: ", to_string(d->scaling_mode)); + + if (mxfp8) { +#if CUBLASMP_VERSION < 801 + NVTE_ERROR("MXFP8 GEMM requires cuBLASMp 0.8.1+."); +#else + NVTE_CHECK(a->with_gemm_swizzled_scales, + "MXFP8 scales for A are not in format expected by GEMM"); + NVTE_CHECK(b->with_gemm_swizzled_scales, + "MXFP8 scales for B are not in format expected by GEMM"); +#endif } - // Mirror cublaslt_gemm.cu's CanonicalizeGemmInput for FP8 tensor scaling: depending on the - // architecture and the quantizer's usage modes, the appropriate data + scale_inv - // may live on the rowwise or columnwise side of the tensor. - // * Hopper (!nvte_is_non_tn_fp8_gemm_supported): only TN FP8 GEMMs are supported, so an - // FP8 input not already in TN orientation must be swapped to its columnwise (transposed) - // view and the transpose flag flipped. - // * Blackwell+ (nvte_is_non_tn_fp8_gemm_supported): any FP8 GEMM layout is supported, but - // the quantizer usage may have only been set to columnwise. In that case, fall back to - // the columnwise view and flip the transpose flag so the GEMM sees the matching data and - // scale_inv pair. - // The original tensor is never modified; a new Tensor view aliases the columnwise pointers. + // Mirror cublaslt_gemm.cu's input canonicalization. Tensor FP8 columnwise data is a + // transposed view, while MXFP8 columnwise data keeps the logical shape. const bool fp8_needs_tn = !nvte_is_non_tn_fp8_gemm_supported(); - auto canonicalize_fp8_input = [fp8_needs_tn](const Tensor* t, bool current_trans, bool want_trans, - const char* side) -> std::pair { + auto canonicalize_input = [fp8_needs_tn](const Tensor* t, bool current_trans, bool is_a, + const char* side) -> std::pair { + auto use_columnwise = [t](bool new_trans) -> std::pair { + Tensor view = *t; + view.data = t->columnwise_data; + view.scale_inv = t->columnwise_scale_inv; + view.amax = t->columnwise_amax; + return {view, new_trans}; + }; + + if (is_mxfp_scaling(t->scaling_mode)) { + if (is_a) { + if (current_trans) { + NVTE_CHECK(t->has_data(), "MXFP8 transposed input A is missing row-wise data"); + return {*t, current_trans}; + } + NVTE_CHECK(t->has_columnwise_data(), + "MXFP8 non-transposed input A is missing column-wise data"); + return use_columnwise(current_trans); + } + if (current_trans) { + NVTE_CHECK(t->has_columnwise_data(), + "MXFP8 transposed input B is missing column-wise data"); + return use_columnwise(current_trans); + } + NVTE_CHECK(t->has_data(), "MXFP8 non-transposed input B is missing row-wise data"); + return {*t, current_trans}; + } + if (!is_fp8_dtype(t->dtype())) { return {*t, current_trans}; } + + const bool want_trans = is_a; const bool hopper_tn_swap = fp8_needs_tn && current_trans != want_trans; const bool blackwell_missing_rowwise = !fp8_needs_tn && !t->has_data(); if (!hopper_tn_swap && !blackwell_missing_rowwise) { @@ -289,16 +324,11 @@ void cublasmp_gemm(InitMatricesFn init_matrices_fn, NVTECommGemmCtx* ctx, NVTECo } NVTE_CHECK(t->has_columnwise_data() && is_fp8_dtype(t->columnwise_data.dtype), "cuBLASMp FP8 GEMM input ", side, " is missing column-wise usage"); - Tensor view; - view.scaling_mode = t->scaling_mode; - view.data = t->columnwise_data; - view.scale_inv = t->columnwise_scale_inv; - // Columnwise data is the transposed view of the original — flip the transpose flag. - return {view, !current_trans}; + return use_columnwise(!current_trans); }; - auto [a_used, transa_eff] = canonicalize_fp8_input(a, transa, /*want_trans=*/true, "A"); - auto [b_used, transb_eff] = canonicalize_fp8_input(b, transb, /*want_trans=*/false, "B"); + auto [a_used, transa_eff] = canonicalize_input(a, transa, /*is_a=*/true, "A"); + auto [b_used, transb_eff] = canonicalize_input(b, transb, /*is_a=*/false, "B"); transa = transa_eff; transb = transb_eff; @@ -321,20 +351,30 @@ void cublasmp_gemm(InitMatricesFn init_matrices_fn, NVTECommGemmCtx* ctx, NVTECo sizeof algo_attr)); const cublasMpMatmulMatrixScale_t scale_mode = CUBLASMP_MATMUL_MATRIX_SCALE_SCALAR_FP32; - if (is_fp8_dtype(a_used.dtype())) { - NVTE_CHECK(a_used.scale_inv.dptr, "Scaling must be set for FP8 dtype"); + auto get_input_scale_mode = [&](const Tensor& t) { +#if CUBLASMP_VERSION >= 801 + if (is_mxfp_scaling(t.scaling_mode)) { + return CUBLASMP_MATMUL_MATRIX_SCALE_VEC32_UE8M0; + } +#endif + return scale_mode; + }; + if (is_fp8_dtype(a_used.dtype()) || is_mxfp_scaling(a_used.scaling_mode)) { + const cublasMpMatmulMatrixScale_t input_scale_mode = get_input_scale_mode(a_used); + NVTE_CHECK(a_used.scale_inv.dptr, "Scaling must be set for input A"); NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( - ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_A_SCALE_MODE, &scale_mode, - sizeof scale_mode)); + ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_A_SCALE_MODE, + &input_scale_mode, sizeof input_scale_mode)); NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_A_SCALE_POINTER, &a_used.scale_inv.dptr, sizeof(void*))); } - if (is_fp8_dtype(b_used.dtype())) { - NVTE_CHECK(b_used.scale_inv.dptr, "Scaling must be set for FP8 dtype"); + if (is_fp8_dtype(b_used.dtype()) || is_mxfp_scaling(b_used.scaling_mode)) { + const cublasMpMatmulMatrixScale_t input_scale_mode = get_input_scale_mode(b_used); + NVTE_CHECK(b_used.scale_inv.dptr, "Scaling must be set for input B"); NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( - ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_B_SCALE_MODE, &scale_mode, - sizeof scale_mode)); + ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_B_SCALE_MODE, + &input_scale_mode, sizeof input_scale_mode)); NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_B_SCALE_POINTER, &b_used.scale_inv.dptr, sizeof(void*))); diff --git a/transformer_engine/jax/cpp_extensions/ep.py b/transformer_engine/jax/cpp_extensions/ep.py index ca70ea145c8..9e2cf2b67c5 100644 --- a/transformer_engine/jax/cpp_extensions/ep.py +++ b/transformer_engine/jax/cpp_extensions/ep.py @@ -15,17 +15,19 @@ """ import functools +import os from dataclasses import dataclass import jax import jax.numpy as jnp +import numpy as np from jax import dtypes, ffi from jax.sharding import NamedSharding, PartitionSpec 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 +from ..version_utils import is_collective_stream_supported, is_xla_ffi_collectives_supported def _on_collective_stream(func): @@ -69,6 +71,35 @@ def wrapper(*args, **kwargs): # ── Module-level EP config ────────────────────────────────────────────────── +@functools.lru_cache(maxsize=None) +def is_ep_borrowed_comm_built() -> bool: + """Return True if transformer_engine_jax was compiled with the borrowed-comm FFI.""" + try: + return "te_ep_bootstrap_borrowed_comm_ffi" in transformer_engine_jax.registrations() + except Exception: # pylint: disable=broad-except + return False + + +def use_nccl_comm_from_xla() -> bool: + """Return True when EP should borrow XLA's NCCL comm instead of self-hosting NCCL. + + Auto-selected when both the build and the installed JAX support the XLA + collectives FFI extension. NVTE_JAX_EP_NCCL_COMM_FROM_XLA=1/0 is an internal + override for tests, not a supported user knob. + """ + env = os.environ.get("NVTE_JAX_EP_NCCL_COMM_FROM_XLA") + if env is not None: + forced_on = env not in ("0", "", "false", "False") + if forced_on and not is_ep_borrowed_comm_built(): + raise RuntimeError( + "NVTE_JAX_EP_NCCL_COMM_FROM_XLA is set but transformer_engine_jax was built " + "without the EP borrowed-comm path (XLA collectives FFI headers were " + "unavailable at build time). Unset it to use the self-hosted NCCL comm." + ) + return forced_on + return is_ep_borrowed_comm_built() and is_xla_ffi_collectives_supported() + + @dataclass(frozen=True) class EpConfig: """Snapshot of the EP bootstrap config (see ep_bootstrap). @@ -91,6 +122,39 @@ class EpConfig: _ep_config: EpConfig = None +# Fixed sentinel keeps EP on its own private comm so it never aliases an XLA +# collective over the same devices. Must stay in [0, 2**63 - 1]. +# 0x54454550 spells "TEEP". +EP_COMMUNICATION_ID = 0x54454550 + + +def run_borrowed_comm_bootstrap( + mesh, replica_groups_flat, group_size, communication_id=EP_COMMUNICATION_ID +): + """Initialize EPBackend on the borrowed XLA comm (one-shot, all devices).""" + try: + from jax import shard_map # top-level since v0.8.0 + except ImportError: # older JAX + from jax.experimental.shard_map import shard_map + + all_axes = tuple(mesh.axis_names) + spec = PartitionSpec(all_axes) + world = int(np.prod([mesh.shape[a] for a in all_axes])) + rg = np.asarray(replica_groups_flat, np.int64) + gs = np.int64(group_size) + cid = np.int64(communication_id) + + def _body(x): + out_type = jax.ShapeDtypeStruct(x.shape, x.dtype) + return ffi.ffi_call("te_ep_bootstrap_borrowed_comm_ffi", out_type, has_side_effect=True)( + x, replica_groups=rg, group_size=gs, communication_id=cid + ) + + dummy = jnp.zeros((world,), dtype=jnp.uint8) + fn = jax.jit(shard_map(_body, mesh=mesh, in_specs=spec, out_specs=spec)) + jax.block_until_ready(fn(dummy)) + + def set_ep_config(config: EpConfig) -> None: """Cache the EP config for abstract-eval / sharding helpers. Call once.""" global _ep_config diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index a7bf8d21d55..ee486585c4d 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -207,8 +207,11 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedMoEAuxLossBackwardHandler); 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, - bool drop_on_overflow); + bool drop_on_overflow, bool borrowed_comm); void ReleaseEpResources(); +// Atexit-safe variant of ReleaseEpResources; never shuts down a borrowed +// backend (see definition). +void ReleaseEpResourcesAtExit(); // Return the handle_mem byte size for a layer config. size_t EpHandleMemSize(int top_k, size_t dispatch_output_per_expert_alignment); @@ -224,6 +227,11 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(EpCombineHandler); XLA_FFI_DECLARE_HANDLER_SYMBOL(EpDispatchBwdHandler); XLA_FFI_DECLARE_HANDLER_SYMBOL(EpCombineBwdHandler); +// EP-specific execute stage of the borrowed-comm bootstrap op (see +// tex.ep.use_nccl_comm_from_xla). The prepare stage is the generic +// FfiRequestCliqueHandler in extensions/ffi_collectives.h. +XLA_FFI_DECLARE_HANDLER_SYMBOL(EpBootstrapBorrowedCommHandler); + // TopK XLA_FFI_DECLARE_HANDLER_SYMBOL(TopkHandler); pybind11::tuple GetTopkWorkspaceSizes(int batch_size, int seq_len, int k); diff --git a/transformer_engine/jax/csrc/extensions/ep.cpp b/transformer_engine/jax/csrc/extensions/ep.cpp index aa5ed27faa4..67bc7c5f31b 100644 --- a/transformer_engine/jax/csrc/extensions/ep.cpp +++ b/transformer_engine/jax/csrc/extensions/ep.cpp @@ -18,6 +18,7 @@ #include "../extensions.h" #include "common.h" +#include "ffi_collectives.h" #include "transformer_engine/gemm.h" namespace transformer_engine { @@ -36,26 +37,32 @@ struct EpBootstrapParams { int max_num_sms = 0; NVTEDType max_token_dtype = kNVTEBFloat16; bool drop_on_overflow = false; + // When set, EP borrows XLA's comm (see below): no ncclCommInitRank at + // bootstrap; nvte_ep_initialize is deferred to the first executable that + // fetches the borrowed communicator. + bool borrowed_comm = false; }; +static NVTEEpGroupConfig MakeEpGroupConfig(const EpBootstrapParams& p) { + return NVTEEpGroupConfig{.struct_size = sizeof(NVTEEpGroupConfig), + .ep_size = p.ep_size, + .num_experts = p.num_experts, + .max_tokens_per_rank = p.max_tokens_per_rank, + .max_recv_tokens_per_rank = p.max_recv_tokens_per_rank, + .hidden_dim = p.hidden_dim, + .num_comm_sms = p.max_num_sms, + .max_token_dtype = p.max_token_dtype, + .zero_copy = 0, + .drop_on_overflow = p.drop_on_overflow}; +} + class EpResources { public: explicit EpResources(const EpBootstrapParams& p) { ncclUniqueId uid; std::memcpy(&uid, p.uid_bytes.data(), sizeof(uid)); NVTE_CHECK_NCCL(ncclCommInitRank(&comm_, p.ep_size, uid, p.rank_within_group)); - // zero_copy=0: JAX EP path always stages payloads; the zero-copy fast path - // requires NVTECommWindow-backed tensors, which JAX bindings don't expose. - NVTEEpGroupConfig cfg{.struct_size = sizeof(NVTEEpGroupConfig), - .ep_size = p.ep_size, - .num_experts = p.num_experts, - .max_tokens_per_rank = p.max_tokens_per_rank, - .max_recv_tokens_per_rank = p.max_recv_tokens_per_rank, - .hidden_dim = p.hidden_dim, - .num_comm_sms = p.max_num_sms, - .max_token_dtype = p.max_token_dtype, - .zero_copy = 0, - .drop_on_overflow = p.drop_on_overflow}; + NVTEEpGroupConfig cfg = MakeEpGroupConfig(p); try { nvte_ep_initialize(static_cast(comm_), &cfg); } catch (...) { @@ -97,6 +104,21 @@ bool g_ep_params_set = false; std::weak_ptr g_ep_resources_weak; // Python-held anchor so trace-time handle_mem allocs find EPBackend ready. std::shared_ptr g_ep_resources_anchor; +// Borrowed-comm path: EPBackend is initialized once from a borrowed communicator. +bool g_ep_xla_initialized = false; + +#ifdef XLA_FFI_COLLECTIVES_AVAILABLE +// Idempotently initialize EPBackend on a borrowed communicator. Safe to call +// from every executable that fetches the comm; only the first call initializes. +void EnsureEpBackendFromBorrowedComm(ncclComm_t comm) { + std::lock_guard lock(g_ep_mu); + if (g_ep_xla_initialized) return; + NVTE_CHECK(g_ep_params_set, "EP bootstrap params not set before borrowing XLA comm."); + NVTEEpGroupConfig cfg = MakeEpGroupConfig(g_ep_params); + nvte_ep_initialize(static_cast(comm), &cfg); + g_ep_xla_initialized = true; +} +#endif // collectives header available std::shared_ptr AcquireEpResources() { std::lock_guard lock(g_ep_mu); @@ -128,14 +150,14 @@ struct EpConfig { 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, - bool drop_on_overflow) { + bool drop_on_overflow, bool borrowed_comm) { 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)."); std::shared_ptr anchor; { std::lock_guard lock(g_ep_mu); - NVTE_CHECK(!g_ep_resources_anchor, + NVTE_CHECK(!g_ep_resources_anchor && !g_ep_xla_initialized, "EP bootstrap already initialized; call release_ep_resources() before re-init."); std::memcpy(g_ep_params.uid_bytes.data(), uid_str.data(), 128); g_ep_params.ep_size = ep_size; @@ -147,8 +169,11 @@ void SetEpBootstrapParams(pybind11::bytes unique_id_bytes_obj, int ep_size, int 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.borrowed_comm = borrowed_comm; g_ep_params_set = true; } + // Borrowed-comm path defers NCCL init to the first executable; nothing eager. + if (borrowed_comm) return; // Acquire outside the lock: EpResources ctor runs ncclCommInitRank which is // a collective and may block on peer ranks. anchor = AcquireEpResources(); @@ -157,7 +182,30 @@ void SetEpBootstrapParams(pybind11::bytes unique_id_bytes_obj, int ep_size, int } // Drops the anchor; comm tears down once the last executable also releases. +// For the borrowed-comm path, tears down EPBackend while the borrowed comm is +// still alive (call from ep_finalize, not atexit). void ReleaseEpResources() { + std::shared_ptr to_drop; + { + std::lock_guard lock(g_ep_mu); + to_drop = std::move(g_ep_resources_anchor); + if (g_ep_xla_initialized) { + nvte_ep_shutdown(); + g_ep_xla_initialized = false; + } + } + // to_drop dtor runs outside the lock. +} + +// Atexit-only variant: drops the self-hosted anchor (its destructor safely +// tears down its own NCCL comm while the process is still alive) but never +// touches a borrowed backend. Ep bootstrap/finalize can switch comm paths +// mid-process (e.g. tests), so by the time atexit fires the live backend may +// be borrowed-comm; XLA may already be tearing down that comm, and calling +// nvte_ep_shutdown() on it here would race that teardown. A borrowed backend +// left live at process exit is handled separately by EPBackend's own +// atexit-safe static destructor (skips NCCL calls). +void ReleaseEpResourcesAtExit() { std::shared_ptr to_drop; { std::lock_guard lock(g_ep_mu); @@ -186,11 +234,16 @@ pybind11::capsule GetEpInstanceStateTypeInfoCapsule() { static ::xla::ffi::ErrorOr> EpInstantiateImpl() { auto state = std::make_unique(); + { + // Borrowed-comm path: XLA owns the comm, so there is nothing self-hosted to + // acquire or pin per-executable; EPBackend is initialized by the bootstrap op. + std::lock_guard lock(g_ep_mu); + if (g_ep_params.borrowed_comm) return state; + } try { state->resources = AcquireEpResources(); } catch (const std::exception& e) { - return ::xla::ffi::Unexpected( - ::xla::ffi::Error::Internal(std::string("EP instantiate failed: ") + e.what())); + return ::xla::ffi::Unexpected(ffi_internal_error("EP instantiate failed: ", e)); } return state; } @@ -479,6 +532,57 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(EpCombineBwdHandler, EpCombineBwdFFI, .Attrs(), FFI_CudaGraph_Traits); +// -- Borrowed-comm path ------------------------------------------------------- +// +// Instead of a self-hosted ncclCommInitRank, EPBackend borrows the communicator +// XLA already owns for the EP replica groups. A one-shot bootstrap op fetches it +// and initializes the backend once; the per-step ops then share the same FFI +// targets as the self-hosted path (see EpInstantiateImpl). Auto-selected from +// Python (see tex.ep.use_nccl_comm_from_xla). The prepare stage is the generic +// FfiRequestCliqueHandler (see ffi_collectives.cpp); only the EP-specific +// execute stage lives here. +#ifdef XLA_FFI_COLLECTIVES_AVAILABLE + +// Execute stage: fetch the borrowed comm and initialize EPBackend once. The +// token buffer is copied straight through only to give the op an input->output +// data dependency, so it stays ordered on the borrowed comm's stream. +Error_Type EpBootstrapBorrowedCommFFI(cudaStream_t stream, EpInstanceState* ep_state, + FfiCollectivesCtx coll, Buffer_Type token, Result_Type out, + Span_Type replica_groups, int64_t group_size, + int64_t communication_id) { + (void)ep_state; + auto groups = ffi_collectives::ReplicaGroupsFromFlat(replica_groups.begin(), + replica_groups.size(), group_size); + auto comm_or = ffi_collectives::GetComm(coll, groups, communication_id); + if (comm_or.has_error()) return comm_or.error(); + ncclComm_t comm = comm_or.value(); + NVTE_CHECK(comm != nullptr, "XLA returned a null EP communicator."); + try { + EnsureEpBackendFromBorrowedComm(comm); + } catch (const std::exception& e) { + return ffi_internal_error("EP borrowed-comm bootstrap failed: ", e); + } + const size_t bytes = token.size_bytes(); + if (bytes > 0) { + NVTE_CHECK_CUDA(cudaMemcpyAsync(out->untyped_data(), token.untyped_data(), bytes, + cudaMemcpyDeviceToDevice, stream)); + } + return ffi_with_cuda_error_check(); +} + +XLA_FFI_DEFINE_HANDLER_SYMBOL(EpBootstrapBorrowedCommHandler, EpBootstrapBorrowedCommFFI, + FFI::Bind() + .Ctx() // stream + .Ctx<::xla::ffi::State>() // EP state + .Ctx<::xla::ffi::Extension>() + .Arg() // token (identity in) + .Ret() // token (identity out) + .Attr>("replica_groups") + .Attr("group_size") + .Attr("communication_id")); + +#endif // collectives header available + } // namespace jax } // namespace transformer_engine diff --git a/transformer_engine/jax/csrc/extensions/ffi.cpp b/transformer_engine/jax/csrc/extensions/ffi.cpp index 6bb2f182342..f9500014e62 100644 --- a/transformer_engine/jax/csrc/extensions/ffi.cpp +++ b/transformer_engine/jax/csrc/extensions/ffi.cpp @@ -11,7 +11,7 @@ namespace transformer_engine { namespace jax { // For XLA_FFI_DataType Enum Reference: https://github.com/openxla/xla/blob/d054e8366c4e8807726961feeb28b1cdba681888/xla/ffi/api/c_api.h#L163-L186 -DType convert_ffi_datatype_to_te_dtype(const xla::ffi::DataType &type) { +DType convert_ffi_datatype_to_te_dtype(const xla::ffi::DataType& type) { switch (type) { // Using this for E8M0 case xla::ffi::DataType::U8: @@ -61,5 +61,9 @@ Error_Type ffi_with_cuda_error_check() { return Error_Type::Success(); } +Error_Type ffi_internal_error(const std::string& context, const std::exception& e) { + return Error_Type::Internal(context + e.what()); +} + } // namespace jax } // namespace transformer_engine diff --git a/transformer_engine/jax/csrc/extensions/ffi.h b/transformer_engine/jax/csrc/extensions/ffi.h index f9d327102bb..8543ede8548 100644 --- a/transformer_engine/jax/csrc/extensions/ffi.h +++ b/transformer_engine/jax/csrc/extensions/ffi.h @@ -8,6 +8,7 @@ #include #include +#include #include "common/util/logging.h" @@ -19,6 +20,8 @@ using Result_Type = xla::ffi::Result; using Variadic_Buffer_Type = xla::ffi::RemainingArgs; using Variadic_Result_Type = xla::ffi::RemainingRets; using Error_Type = xla::ffi::Error; +template +using Span_Type = xla::ffi::Span; using FFI = xla::ffi::Ffi; using FFI_Stream_Type = xla::ffi::PlatformStream; using Dictionary = xla::ffi::Dictionary; @@ -31,6 +34,9 @@ DType convert_ffi_datatype_to_te_dtype(const xla::ffi::DataType& type); Error_Type ffi_with_cuda_error_check(); +// Wraps a caught std::exception as an Error_Type, prefixed with `context`. +Error_Type ffi_internal_error(const std::string& context, const std::exception& e); + // source_location is not available in C++17, so we implement it ourselves #if defined(__GNUC__) || defined(__clang__) #define CURRENT_FILE __builtin_FILE() diff --git a/transformer_engine/jax/csrc/extensions/ffi_collectives.cpp b/transformer_engine/jax/csrc/extensions/ffi_collectives.cpp new file mode 100644 index 00000000000..3e1323e535f --- /dev/null +++ b/transformer_engine/jax/csrc/extensions/ffi_collectives.cpp @@ -0,0 +1,106 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "ffi_collectives.h" + +#ifdef XLA_FFI_COLLECTIVES_AVAILABLE + +#include + +#include "ffi.h" + +namespace transformer_engine { +namespace jax { + +namespace ffi_collectives { + +namespace { + +::xla::ffi::Error TakeError(const XLA_FFI_Api* api, XLA_FFI_Error* err) { + std::string msg = ::xla::ffi::internal::GetErrorMessage(api, err); + ::xla::ffi::internal::DestroyError(api, err); + return ::xla::ffi::Error::Internal(msg); +} + +} // namespace + +std::vector ToRawGroups(const std::vector>& groups) { + std::vector raw; + raw.reserve(groups.size()); + for (const auto& g : groups) { + raw.push_back(XLA_FFI_ReplicaGroup{g.data(), g.size()}); + } + return raw; +} + +::xla::ffi::Error RequestClique(const FfiCollectivesCtx& ctx, + const std::vector>& groups, + int64_t communication_id) { + std::vector raw = ToRawGroups(groups); + XLA_FFI_Communicator_Request_Args args = {}; + args.struct_size = XLA_FFI_Communicator_Request_Args_STRUCT_SIZE; + args.extension_start = nullptr; + args.group_mode = XLA_FFI_GROUP_FLATTENED_ID; + args.groups = raw.data(); + args.num_groups = raw.size(); + args.communication_id = communication_id; + if (XLA_FFI_Error* err = ctx.ext->request_communicator(ctx.ext, &args)) { + return TakeError(ctx.api, err); + } + return ::xla::ffi::Error::Success(); +} + +::xla::ffi::ErrorOr GetComm(const FfiCollectivesCtx& ctx, + const std::vector>& groups, + int64_t communication_id) { + std::vector raw = ToRawGroups(groups); + XLA_FFI_Communicator_Get_Args args = {}; + args.struct_size = XLA_FFI_Communicator_Get_Args_STRUCT_SIZE; + args.extension_start = nullptr; + args.group_mode = XLA_FFI_GROUP_FLATTENED_ID; + args.groups = raw.data(); + args.num_groups = raw.size(); + args.communication_id = communication_id; + args.communicator = nullptr; + if (XLA_FFI_Error* err = ctx.ext->get_communicator(ctx.ext, &args)) { + return TakeError(ctx.api, err); + } + return reinterpret_cast(args.communicator); +} + +std::vector> ReplicaGroupsFromFlat(const int64_t* flat, size_t count, + int64_t group_size) { + NVTE_CHECK(group_size > 0, "ReplicaGroupsFromFlat: group_size must be > 0, got ", group_size); + NVTE_CHECK(count % static_cast(group_size) == 0, + "ReplicaGroupsFromFlat: flat replica-group buffer size (", count, + ") must be a multiple of group_size (", group_size, ")"); + std::vector> groups; + for (size_t off = 0; off + group_size <= count; off += group_size) { + groups.emplace_back(flat + off, flat + off + group_size); + } + return groups; +} + +} // namespace ffi_collectives + +Error_Type FfiRequestCliqueFFI(FfiCollectivesCtx coll, Span_Type replica_groups, + int64_t group_size, int64_t communication_id) { + auto groups = ffi_collectives::ReplicaGroupsFromFlat(replica_groups.begin(), + replica_groups.size(), group_size); + return ffi_collectives::RequestClique(coll, groups, communication_id); +} + +XLA_FFI_DEFINE_HANDLER_SYMBOL(FfiRequestCliqueHandler, FfiRequestCliqueFFI, + FFI::BindPrepare() + .Ctx<::xla::ffi::Extension>() + .Attr>("replica_groups") + .Attr("group_size") + .Attr("communication_id")); + +} // namespace jax +} // namespace transformer_engine + +#endif // collectives header available diff --git a/transformer_engine/jax/csrc/extensions/ffi_collectives.h b/transformer_engine/jax/csrc/extensions/ffi_collectives.h new file mode 100644 index 00000000000..3063a8b3dec --- /dev/null +++ b/transformer_engine/jax/csrc/extensions/ffi_collectives.h @@ -0,0 +1,95 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file ffi_collectives.h + * \brief Borrow the XLA-owned NCCL communicator inside any FFI handler. + * + * Not EP-specific -- any FFI handler that wants XLA's comm can request the + * clique (prepare stage) and fetch the borrowed ncclComm_t (execute stage). + * Absent on older XLA, in which case the borrow path must not be selected. + */ + +#ifndef TRANSFORMER_ENGINE_JAX_CSRC_EXTENSIONS_FFI_COLLECTIVES_H_ +#define TRANSFORMER_ENGINE_JAX_CSRC_EXTENSIONS_FFI_COLLECTIVES_H_ + +// XLA_FFI_COLLECTIVES_AVAILABLE is the single source of truth for "is the +// XLA collectives FFI extension available"; callers add their own build gates +// (e.g. NVTE_WITH_NCCL_EP) on top of this. +#if __has_include("xla/ffi/api/collectives_c_api.h") +#define XLA_FFI_COLLECTIVES_AVAILABLE 1 + +#include + +#include +#include + +#include "xla/ffi/api/collectives_c_api.h" +#include "xla/ffi/api/ffi.h" + +namespace transformer_engine { +namespace jax { + +// Decoded context: the FFI api table plus the found collectives extension. +struct FfiCollectivesCtx { + const XLA_FFI_Api* api = nullptr; + const XLA_FFI_Collectives_Extension* ext = nullptr; +}; + +// Trait type for ::xla::ffi::Extension. The public FFI +// CtxDecoding looks the extension up by kExtensionType and hands us a typed +// context, so we do not depend on XLA-internal headers that jaxlib omits. +// +// TODO: XLA ships an equivalent trait + wrapper (xla::ffi::Collectives / +// xla::ffi::Communicator) in xla/ffi/api/collectives_ffi.h. Drop this trait +// and RequestClique/GetComm/ToRawGroups below in favor of that once both: +// 1. jaxlib packages collectives_ffi.h (not shipped as of jaxlib +// 0.11.2.dev20260913), and +// 2. XLA's CollectivesExtensionBase gains a Support() override accepting +// any minor within the same major (currently exact-match only), so we +// do not regress the forward-compat policy below. +struct FfiCollectives { + using Type = FfiCollectivesCtx; + using CExtension = XLA_FFI_Collectives_Extension; + static constexpr const char* kName = "CollectivesExtension"; + static constexpr int32_t kExtensionType = XLA_FFI_Extension_Collectives; + static constexpr int32_t kMajorVersion = XLA_FFI_Extension_Collectives_MajorVersion; + static constexpr int32_t kMinorVersion = XLA_FFI_Extension_Collectives_MinorVersion; + // Accept any minor within the same major so a newer runtime still binds. + static bool Support(int32_t major, int32_t /*minor*/) { return major == kMajorVersion; } + static Type Create(const XLA_FFI_Api* api, const CExtension* ext) { return Type{api, ext}; } +}; + +namespace ffi_collectives { + +std::vector ToRawGroups(const std::vector>& groups); + +// Prepare stage: ask XLA to acquire the clique for `groups` (flattened-id mode). +::xla::ffi::Error RequestClique(const FfiCollectivesCtx& ctx, + const std::vector>& groups, + int64_t communication_id); + +// Execute stage: fetch the borrowed communicator (ncclComm_t on XLA:GPU). +::xla::ffi::ErrorOr GetComm(const FfiCollectivesCtx& ctx, + const std::vector>& groups, + int64_t communication_id); + +// Rebuild ragged replica groups from a flat buffer of equal-size groups. +std::vector> ReplicaGroupsFromFlat(const int64_t* flat, size_t count, + int64_t group_size); + +} // namespace ffi_collectives + +// Generic prepare-stage handler: any FFI op that borrows XLA's comm binds this +// to request the clique before execute. Reads int64 attrs "replica_groups" +// (flat, equal-size groups), "group_size", and "communication_id". +XLA_FFI_DECLARE_HANDLER_SYMBOL(FfiRequestCliqueHandler); + +} // namespace jax +} // namespace transformer_engine + +#endif // collectives header available + +#endif // TRANSFORMER_ENGINE_JAX_CSRC_EXTENSIONS_FFI_COLLECTIVES_H_ diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index e65cea9fa0e..5538e3538f2 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -7,6 +7,7 @@ #include "../extensions.h" #include "cgemm_helper.h" #include "common/util/cuda_runtime.h" +#include "ffi_collectives.h" // FfiRequestCliqueHandler (borrowed-comm prepare) #include "transformer_engine/gemm.h" namespace transformer_engine { @@ -124,6 +125,17 @@ pybind11::dict Registrations() { dict["te_ep_combine_bwd_ffi"] = pybind11::dict(pybind11::arg("instantiate") = EncapsulateFFI(EpInstantiateHandler), pybind11::arg("execute") = EncapsulateFFI(EpCombineBwdHandler)); + + // Borrowed-comm bootstrap: a one-shot op requests the collective clique + // (generic prepare) and initializes EPBackend from the borrowed comm (EP + // execute). Registered only when the XLA collectives headers were available + // at build time; its absence is how Python detects an unbuilt path. +#ifdef XLA_FFI_COLLECTIVES_AVAILABLE + dict["te_ep_bootstrap_borrowed_comm_ffi"] = + pybind11::dict(pybind11::arg("instantiate") = EncapsulateFFI(EpInstantiateHandler), + pybind11::arg("prepare") = EncapsulateFFI(FfiRequestCliqueHandler), + pybind11::arg("execute") = EncapsulateFFI(EpBootstrapBorrowedCommHandler)); +#endif // collectives header available #endif // NVTE_WITH_NCCL_EP // TopK @@ -161,8 +173,9 @@ PYBIND11_MODULE(transformer_engine_jax, m) { 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("drop_on_overflow")); + pybind11::arg("drop_on_overflow"), pybind11::arg("borrowed_comm") = false); m.def("release_ep_resources", &ReleaseEpResources); + m.def("release_ep_resources_at_exit", &ReleaseEpResourcesAtExit); m.def("ep_handle_mem_size", &EpHandleMemSize, pybind11::arg("top_k"), pybind11::arg("dispatch_output_per_expert_alignment") = 0); m.def("get_ep_instance_state_type_id", &GetEpInstanceStateTypeIdCapsule); diff --git a/transformer_engine/jax/ep.py b/transformer_engine/jax/ep.py index 2222a41e482..57aed7cab36 100644 --- a/transformer_engine/jax/ep.py +++ b/transformer_engine/jax/ep.py @@ -100,6 +100,21 @@ def device_to_rank(d): return int(grid[row, 0]), col, int(grid.shape[0]) +def _ep_flattened_replica_groups(mesh, ep_resource): + """FLATTENED_ID replica groups for the EP axis, as a flat int64 array. + + Each group fixes all non-ep mesh coordinates and varies ep. Returns + ``(flat_groups, ep_size)``. + """ + shape = tuple(mesh.shape[a] for a in mesh.axis_names) + ep_pos = mesh.axis_names.index(ep_resource) + ep_size = shape[ep_pos] + world = int(np.prod(shape)) + grid = np.arange(world, dtype=np.int64).reshape(shape) + groups = np.moveaxis(grid, ep_pos, -1).reshape(-1, ep_size) + return groups.reshape(-1), ep_size + + def ep_bootstrap( world_size, rank, @@ -116,11 +131,14 @@ def ep_bootstrap( 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. Axes orthogonal to EP (tp, pp, - cp, ...) are supported and replicated across EP tensors. + cp, ...) are supported and replicated across EP tensors. Auto-selects + between self-hosted NCCL and XLA's borrowed comm (see + tex.ep.use_nccl_comm_from_xla). Args: world_size: Total number of processes (product of all mesh axes). - rank: Global rank of the calling process. + rank: Global rank of the calling process. Unused on the borrowed-comm + path, where each device's identity comes from its mesh position. num_experts: Total experts across the EP group. max_tokens_per_rank: Max tokens one rank dispatches per step (sizes send buffers). recv_capacity_per_rank: Max tokens one rank receives per step; set to @@ -175,6 +193,42 @@ def ep_bootstrap( if num_experts % ep_size != 0: raise ValueError(f"num_experts ({num_experts}) must be divisible by ep_size ({ep_size}).") + common_cfg = { + "world_size": world_size, + "rank": rank, + "ep_size": ep_size, + "num_ep_groups": num_ep_groups, + "num_experts": num_experts, + "num_local_experts": num_experts // ep_size, + "max_tokens_per_rank": max_tokens_per_rank, + "recv_capacity_per_rank": recv_capacity_per_rank, + "hidden_dim": hidden_dim, + } + + # Borrowed-comm path (auto-selected by tex.ep.use_nccl_comm_from_xla): XLA + # owns the EP communicator, so a one-shot bootstrap op fetches it and + # initializes EPBackend instead of a host-side UID exchange. + if tex.ep.use_nccl_comm_from_xla(): + replica_groups, ep_group_size = _ep_flattened_replica_groups(mesh, ep_resource) + communication_id = tex.ep.EP_COMMUNICATION_ID + transformer_engine_jax.set_ep_bootstrap_params( + bytes(128), + ep_size, + 0, + num_experts, + max_tokens_per_rank, + recv_capacity_per_rank, + 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), + borrowed_comm=True, + ) + tex.ep.set_ep_config(tex.ep.EpConfig(**common_cfg)) + # Initialize EPBackend now so trace-time handle_mem_size finds it ready. + tex.ep.run_borrowed_comm_bootstrap(mesh, replica_groups, ep_group_size, communication_id) + return + UID_SIZE = 128 root_rank, rank_within_group, _num_domains = _ep_domain_for_rank(mesh, ep_resource, rank) is_color_root = rank_within_group == 0 @@ -206,34 +260,28 @@ def ep_bootstrap( ) # Release the C++ anchor at interpreter shutdown so RAII can tear down NCCL. + # release_ep_resources_at_exit (not release_ep_resources): a later + # ep_finalize + borrowed-comm ep_bootstrap in this same process must not + # leave this self-hosted-registered hook shutting down the borrowed + # backend at exit (see ReleaseEpResourcesAtExit). global _atexit_registered if not _atexit_registered: - atexit.register(transformer_engine_jax.release_ep_resources) + atexit.register(transformer_engine_jax.release_ep_resources_at_exit) _atexit_registered = True - tex.ep.set_ep_config( - tex.ep.EpConfig( - world_size=world_size, - rank=rank, - ep_size=ep_size, - num_ep_groups=num_ep_groups, - num_experts=num_experts, - num_local_experts=num_experts // ep_size, - max_tokens_per_rank=max_tokens_per_rank, - recv_capacity_per_rank=recv_capacity_per_rank, - hidden_dim=hidden_dim, - ) - ) + tex.ep.set_ep_config(tex.ep.EpConfig(**common_cfg)) 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. + configs); a normal run just exits without calling this. Self-hosted NCCL + also tears down via atexit; the borrowed-comm path has no atexit hook, so + use this for mid-process teardown there. 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() diff --git a/transformer_engine/jax/version_utils.py b/transformer_engine/jax/version_utils.py index 500c859b4c8..18ee255f8fc 100644 --- a/transformer_engine/jax/version_utils.py +++ b/transformer_engine/jax/version_utils.py @@ -84,6 +84,26 @@ def is_collective_stream_supported() -> bool: return True +# Minimum JAX version whose XLA ships the FFI collectives extension (lets an FFI +# handler fetch XLA's own communicator). Conservative floor: the first version +# this was verified on. +_XLA_FFI_COLLECTIVES_NIGHTLY_FLOOR = "0.11.2.dev20260828" +_XLA_FFI_COLLECTIVES_STABLE_FLOOR = "0.11.2" + + +@lru_cache(maxsize=None) +def is_xla_ffi_collectives_supported() -> bool: + """Return True if the installed JAX exposes the XLA FFI collectives extension. + + Not EP-specific; gates auto-selection of the EP borrowed-comm path (see + cpp_extensions.ep.use_nccl_comm_from_xla). + """ + v = PkgVersion(get_pkg_version("jax")) + if v.dev is not None: + return v >= PkgVersion(_XLA_FFI_COLLECTIVES_NIGHTLY_FLOOR) + return v >= PkgVersion(_XLA_FFI_COLLECTIVES_STABLE_FLOOR) + + def is_triton_extension_supported() -> bool: """Return True if the current JAX version supports Triton kernel dispatch. @@ -98,6 +118,7 @@ def is_triton_extension_supported() -> bool: "jax_version_meet_requirement", "is_triton_autotuned_alias_safe", "is_collective_stream_supported", + "is_xla_ffi_collectives_supported", "is_triton_extension_supported", "TRITON_EXTENSION_MIN_JAX_VERSION", "TRITON_EXTENSION_CUDA_GRAPH_MIN_JAX_VERSION", diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py index 908e88f30d1..98262bd99ee 100644 --- a/transformer_engine/pytorch/module/_common.py +++ b/transformer_engine/pytorch/module/_common.py @@ -12,7 +12,9 @@ from .. import cpp_extensions as tex from ..constants import TE_DType +from ..distributed import in_fp8_activation_recompute_phase from ..export import is_in_onnx_export_mode +from ..quantization import FP8GlobalStateManager from ..tensor.hybrid_tensor import HybridQuantizer from ..utils import get_default_init_method @@ -320,3 +322,17 @@ def assert_empty(self): assert self.enabled is True, "delay_wgrad_compute is not enabled" rank = torch.distributed.get_rank() assert self.context.empty(), f"Queue is not empty. rank {rank}" + + +def check_fp8_reduce_and_update(restore_first_module: bool = False) -> bool: + """Whether this module's backward should reduce and update the FP8 scaling factors. + + Consumes the "first FP8 module" flag, restored when the forward is a + recomputation so the flag survives for the real forward's owner. + """ + qstate = FP8GlobalStateManager.quantization_state + first_fp8_module = qstate.is_first_fp8_module + result = FP8GlobalStateManager.is_first_fp8_module() + if restore_first_module or in_fp8_activation_recompute_phase(): + qstate.is_first_fp8_module = first_fp8_module + return result diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 561e813348c..c50fb2c4672 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -6,7 +6,8 @@ import os import warnings import weakref -from typing import Callable, Dict, Optional, Tuple, Union, List +from dataclasses import dataclass +from typing import Any, Callable, Dict, Optional, Tuple, Union, List from functools import reduce from operator import mul as multiply_op @@ -41,7 +42,6 @@ init_method_constant, nvtx_range_pop, nvtx_range_push, - requires_grad, needs_quantized_gemm, get_nvtx_range_context, ) @@ -52,7 +52,6 @@ symmetric_all_reduce, reduce_scatter_along_first_dim, gather_along_first_dim, - in_fp8_activation_recompute_phase, _fsdp_scatter_tensors, _fsdp_gather_tensors, ) @@ -67,6 +66,7 @@ from ..graph import is_graph_capturing from ._common import ( apply_normalization, + check_fp8_reduce_and_update, noop_cat, set_quantizer_amax_reduction_group, set_quantizer_usage_for_wgrad_all_gather, @@ -79,6 +79,7 @@ prepare_for_saving, restore_from_func_ctx, ) +from ..dynamo import TensorOrQuantized from ...debug.pytorch.debug_state import TEDebugState from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.hybrid_tensor import HybridQuantizer @@ -100,1107 +101,1457 @@ __all__ = ["LayerNormLinear"] -class _LayerNormLinear(torch.autograd.Function): - """LayerNormLinear semi-top level module - Calls custom cuda extensions. - """ +@dataclass(slots=True) +class LayerNormLinearFwdArgs: + """Single-argument bag for the forward path of :class:`_LayerNormLinear`.""" + + # --- Differentiable tensors (also passed positionally to autograd) --- + inp: torch.Tensor + ln_weight: torch.Tensor + ln_bias: Optional[torch.Tensor] + weight: TensorOrQuantized + bias: Optional[torch.Tensor] + + # --- Non-differentiable cached tensors --- + weight_workspace: Optional[TensorOrQuantized] + + # --- requires_grad flags (cached so backward does not re-query) --- + input_requires_grad: bool + ln_weight_requires_grad: bool + ln_bias_requires_grad: bool + weight_requires_grad: bool + bias_requires_grad: bool + + # --- Quantizers --- + input_quantizer: Optional[Quantizer] + weight_quantizer: Optional[Quantizer] + output_quantizer: Optional[Quantizer] + grad_input_quantizer: Optional[Quantizer] + grad_weight_quantizer: Optional[Quantizer] + grad_output_quantizer: Optional[Quantizer] + + # --- Normalization --- + eps: float + normalization: str + zero_centered_gamma: bool + fwd_ln_sm_margin: int + bwd_ln_sm_margin: int + return_layernorm_output: bool + return_layernorm_output_gathered: bool + + # --- Numerical / dtype config --- + activation_dtype: torch.dtype + fp8: bool + fp8_calibration: bool + backward_override: Optional[str] + dgrad_use_split_accumulator: bool + wgrad_use_split_accumulator: bool + debug: bool + + # --- Weight-workspace caching --- + is_first_microbatch: Optional[bool] + cache_weight: bool + skip_fp8_weight_update: Optional[torch.Tensor] + + # --- Tensor / sequence parallelism --- + parallel_mode: Optional[str] + tp_group: Optional[dist_group_type] + tp_size: int + tensor_parallel: bool + sequence_parallel: bool + symmetric_ar_type: Optional[str] + + # --- Userbuffers (comm + GEMM overlap) --- + ub_name: Optional[str] + ub_overlap_ag_fprop: bool + ub_overlap_rs_fprop: bool + ub_overlap_ag_dgrad: bool + ub_overlap_rs_dgrad: bool + ub_bulk_dgrad: bool + ub_bulk_wgrad: bool + + # --- FSDP --- + fsdp_group: Optional[Any] + is_fsdp2: bool + + # --- Weight-grad scheduling --- + fuse_wgrad_accumulation: bool + wgrad_store: Optional[Any] + + # --- Misc --- + cpu_offloading: bool + is_grad_enabled: bool + + def any_requires_grad(self) -> bool: + """Whether any differentiable input requires a gradient.""" + return any( + ( + self.input_requires_grad, + self.ln_weight_requires_grad, + self.ln_bias_requires_grad, + self.weight_requires_grad, + self.bias_requires_grad, + ) + ) - @staticmethod - def forward( - ctx, - inp: torch.Tensor, - ln_weight: torch.Tensor, - ln_bias: Union[torch.Tensor, None], - weight: torch.Tensor, - weight_workspace: Optional[torch.Tensor], - bias: torch.Tensor, - non_tensor_args: Tuple, - ) -> Tuple[torch.Tensor, ...]: - # pylint: disable=missing-function-docstring - - # Reduce number of arguments to autograd function in order - # to reduce CPU overhead due to pytorch arg checking. + +@dataclass(slots=True) +class LayerNormLinearBwdArgs: + """Single-argument bag for the backward path of :class:`_LayerNormLinear`.""" + + # --- Incoming gradients (populated at backward entry) --- + grad_output: Optional[torch.Tensor] = None + grad_ln_out: Optional[torch.Tensor] = None + + # --- Saved / restored tensors (populated at backward entry) --- + inputmat: Optional[torch.Tensor] = None + weight_fp8: Optional[TensorOrQuantized] = None + saved_weight: Optional[TensorOrQuantized] = None + bias: Optional[torch.Tensor] = None + ln_weight: Optional[torch.Tensor] = None + ln_out: Optional[TensorOrQuantized] = None + mu: Optional[torch.Tensor] = None + rsigma: Optional[torch.Tensor] = None + + # --- Quantizers --- + input_quantizer: Optional[Quantizer] = None + weight_quantizer: Optional[Quantizer] = None + grad_input_quantizer: Optional[Quantizer] = None + grad_weight_quantizer: Optional[Quantizer] = None + grad_output_quantizer: Optional[Quantizer] = None + + # --- Differentiability summary --- + use_bias: bool = False + requires_dgrad: bool = False + requires_wgrad: bool = False + ln_out_needs_gather: bool = False + inp_shape: Optional[torch.Size] = None + + # --- Normalization --- + normalization: str = "LayerNorm" + zero_centered_gamma: bool = False + bwd_ln_sm_margin: int = 0 + return_layernorm_output: bool = False + return_layernorm_output_gathered: bool = False + + # --- Numerical / dtype config --- + activation_dtype: Optional[torch.dtype] = None + fp8: bool = False + dgrad_use_split_accumulator: bool = _2X_ACC_DGRAD + wgrad_use_split_accumulator: bool = _2X_ACC_WGRAD + backward_override: Optional[str] = None + is_weight_param_quantized: bool = False + debug: bool = False + + # --- Tensor / sequence parallelism --- + parallel_mode: Optional[str] = None + tp_group: Optional[dist_group_type] = None + tp_size: int = 1 + tensor_parallel: bool = False + sequence_parallel: bool = False + + # --- Userbuffers (comm + GEMM overlap) --- + ub_name: Optional[str] = None + ub_overlap_ag: bool = False + ub_overlap_rs_dgrad: bool = False + ub_bulk_dgrad: bool = False + ub_bulk_wgrad: bool = False + + # --- FSDP --- + fsdp_group: Optional[Any] = None + fsdp_shapes: Any = None + is_fsdp2: bool = False + + # --- Weight-grad scheduling / accumulation --- + is_first_microbatch: Optional[bool] = None + fuse_wgrad_accumulation: bool = False + wgrad_store: Optional[Any] = None + origin_weight_ref: Optional[Any] = None + origin_weight_overwrites_main_grad: bool = False + main_grad_func: Optional[Callable[[], torch.Tensor]] = None + + # --- FP8 reduce-and-update bookkeeping --- + reduce_and_update_bwd_fp8_tensors: bool = False + + # --- Misc --- + cpu_offloading: bool = False + + # --- Per-backward scratch state (populated inside the backward impl) --- + ub_obj_gradout: Optional[Any] = None + + def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None: + """Pull saved tensors from ``ctx`` into the fields backward consumes.""" ( - eps, - is_first_microbatch, - fp8, - fp8_calibration, - wgrad_store, - fuse_wgrad_accumulation, + self.inputmat, + self.weight_fp8, + self.saved_weight, + self.bias, + self.ln_weight, + self.ln_out, + self.mu, + self.rsigma, + ) = restore_from_func_ctx( + ctx + ) # pylint: disable=unbalanced-tuple-unpacking + + +def _layernorm_linear_forward_impl( + args: LayerNormLinearFwdArgs, +) -> Tuple[ + torch.Tensor, + Optional[torch.Tensor], + Optional[torch.Tensor], + Optional[Tuple], + Optional[Dict], +]: + """Forward implementation for the layernorm-linear layer. + + Returns ``(out, ln_out_return, new_weight_workspace, + tensors_to_save_from_forward, ctx_attrs)``. ``new_weight_workspace`` is + the freshly produced FP8 weight workspace (returned alongside ``out`` so + the caller can refresh its cache). The last two are ``None`` when + gradients are disabled. + """ + inp = args.inp + ln_weight = args.ln_weight + ln_bias = args.ln_bias + weight = args.weight + bias = args.bias + input_quantizer = args.input_quantizer + weight_quantizer = args.weight_quantizer + output_quantizer = args.output_quantizer + is_first_microbatch = args.is_first_microbatch + fp8 = args.fp8 + debug = args.debug + cpu_offloading = args.cpu_offloading + tp_group = args.tp_group + tp_size = args.tp_size + sequence_parallel = args.sequence_parallel + activation_dtype = args.activation_dtype + parallel_mode = args.parallel_mode + is_grad_enabled = args.is_grad_enabled + return_layernorm_output = args.return_layernorm_output + return_layernorm_output_gathered = args.return_layernorm_output_gathered + backward_override = args.backward_override + ub_name = args.ub_name + ub_overlap_ag_fprop = args.ub_overlap_ag_fprop + ub_overlap_rs_fprop = args.ub_overlap_rs_fprop + fsdp_group = args.fsdp_group + is_fsdp2 = args.is_fsdp2 + weight_requires_grad = args.weight_requires_grad + + # NVTX label for profiling + nvtx_label = "transformer_engine._LayerNormLinear.forward" + if ub_name is not None: + nvtx_label = f"{nvtx_label}.{ub_name}" + + with_input_all_gather = parallel_mode == "column" and sequence_parallel + + # Make sure input dimensions are compatible + out_features, in_features = weight.shape + inp_shape = inp.shape + assert inp_shape[-1] == in_features, "GEMM not possible" + inp = inp.view((-1, in_features)) + inputmat = inp + if fp8: + assert_dim_for_fp8_exec(inputmat, weight) + + # Cast for native AMP + nvtx_range_push(f"{nvtx_label}.norm_input_cast") + inputmat = cast_if_needed(inputmat, activation_dtype) + ln_weight_cast = cast_if_needed(ln_weight, activation_dtype) + if ln_bias is not None: + ln_bias = cast_if_needed(ln_bias, activation_dtype) + nvtx_range_pop(f"{nvtx_label}.norm_input_cast") + + if is_cpu_offload_enabled(): + start_offload(inputmat) + + tp_world_size = get_distributed_world_size(tp_group) + + backward_needs_input = is_grad_enabled and weight_requires_grad + + # Configure Userbuffers communication (comm+GEMM overlap) + ub_obj = None + ub_type = None + ub_overlap_ag_fprop = ub_overlap_ag_fprop and is_grad_enabled and not return_layernorm_output + if ub_overlap_rs_fprop: + ub_obj = get_ub(ub_name + "_fprop", fp8) + ub_type = tex.CommOverlapType.RS + elif ub_overlap_ag_fprop: + ub_obj = get_ub(ub_name + "_fprop", fp8) + ub_type = tex.CommOverlapType.AG + + # Configure quantizer for norm output + if fp8: + if input_quantizer is None: + raise ValueError("Missing quantizer for input tensor") + input_quantizer.set_usage( + rowwise=True, + columnwise=backward_needs_input and backward_override is None, + ) + if with_input_all_gather and input_quantizer.supports_only_rowwise_all_gather(): + # All-gather is not supported with FP8 column-wise data + input_quantizer.set_usage(columnwise=False) + # Amax reduction group for the input quantizer (column-parallel sequence parallel) + set_quantizer_amax_reduction_group( input_quantizer, - weight_quantizer, - output_quantizer, - grad_input_quantizer, - grad_weight_quantizer, - grad_output_quantizer, - cpu_offloading, - tp_group, - tp_size, - sequence_parallel, - tensor_parallel, - activation_dtype, - parallel_mode, - return_layernorm_output, - return_layernorm_output_gathered, - is_grad_enabled, - fwd_ln_sm_margin, - bwd_ln_sm_margin, - zero_centered_gamma, - normalization, - ub_overlap_ag_fprop, - ub_overlap_rs_fprop, - ub_overlap_ag_dgrad, - ub_overlap_rs_dgrad, - ub_bulk_wgrad, - ub_bulk_dgrad, - ub_name, - fsdp_group, - cache_weight, - skip_fp8_weight_update, - symmetric_ar_type, - debug, - is_fsdp2, - ) = non_tensor_args - if fp8: - backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override + tp_group if (sequence_parallel and parallel_mode == "column") else None, + ) + + # 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 + and not return_layernorm_output + 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 + ) + + # Apply normalization + nvtx_range_push(f"{nvtx_label}.norm") + ln_out, mu, rsigma = apply_normalization( + inputmat, + None, # ln_out + ln_weight_cast, + ln_bias, + args.eps, + input_quantizer if with_quantized_norm else None, + inputmat.dtype, + args.normalization, + args.fwd_ln_sm_margin, + args.zero_centered_gamma, + ) + nvtx_range_pop(f"{nvtx_label}.norm") + + # Store unquantized layer norm output if we need to return it + ln_out_return = None + if return_layernorm_output or return_layernorm_output_gathered: + ln_out_return = ln_out + ln_out_hp = ln_out if backward_override == "high_precision" else None + + # ------------------------------------------------------ + # Prepare GEMM input tensor + # Note: Cast to expected dtype and perform tensor-parallel communication + # ------------------------------------------------------ + nvtx_range_push(f"{nvtx_label}.gemm_input_cast_comm") + ln_out_total = None + if with_input_all_gather: + if return_layernorm_output_gathered: + # Perform all-gather in high precision if gathered + # norm output will be returned + ln_out_total, _ = gather_along_first_dim(ln_out, tp_group) + ln_out_return = ln_out_total + if fp8 or debug: + ln_out = input_quantizer(ln_out) + input_quantizer.set_usage(rowwise=True, columnwise=False) + ln_out_total = input_quantizer(ln_out_total) else: - backward_override = None + quantizer = None + if fp8 or debug: + quantizer = input_quantizer + # custom recipe doesn't need to support quantized AG + if not with_quantized_norm and not custom: + ln_out = quantizer(ln_out) + quantizer.set_usage(rowwise=True, columnwise=False) + if ub_overlap_ag_fprop: # Initialize Userbuffers all-gather + ln_out_total, _ = fill_userbuffers_buffer_for_all_gather( + ub_obj, + ln_out, + quantizer, + tp_group, + ) + else: # Perform NCCL all-gather + ln_out_total, _ = gather_along_first_dim( + ln_out, + tp_group, + quantizer=quantizer, + ) + else: + if (fp8 or debug) and not with_quantized_norm: + ln_out = input_quantizer(ln_out) + ln_out_total = ln_out + nvtx_range_pop(f"{nvtx_label}.gemm_input_cast_comm") + # ------------------------------------------------------ + # GEMM input tensor is ready... + # ------------------------------------------------------ + + # ------------------------------------------------------ + # Prepare weight tensor + # ------------------------------------------------------ + is_dist_weight = is_distributed_weight(args.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 + if fp8 or debug: + is_weight_param_quantized = isinstance(weight, QuantizedTensorStorage) + + # Configure quantizer + # If weight is already quantized, weight._quantizer is its true quantizer. + # for debug mode we create quantizer every iteration, thus we need to set the quantizer states + if is_weight_param_quantized and not debug: + weight_quantizer = weight._quantizer + elif weight_quantizer is not None: + # FSDP2: Skip columnwise/transpose creation during forward + # to avoid accumulating caches across layers. Backward's + # FSDP2 all-gather will recreate them. (Issue #2681) + weight_quantizer.set_usage( + rowwise=True, + columnwise=is_grad_enabled and not is_fsdp2 and backward_override is None, + ) - # NVTX label for profiling - nvtx_label = "transformer_engine._LayerNormLinear.forward" - if ub_name is not None: - nvtx_label = f"{nvtx_label}.{ub_name}" + # Get quantized weight + update_ws = is_first_microbatch is None or is_first_microbatch + weightmat, new_weight_workspace = quantize_weight( + tensor=weight, + quantizer=weight_quantizer, + workspace=args.weight_workspace, + update_workspace=update_ws, + skip_update_flag=args.skip_fp8_weight_update, + fsdp_group=fsdp_group, + workspace_dtype=activation_dtype, + cache=args.cache_weight, + ) - with_input_all_gather = parallel_mode == "column" and sequence_parallel + weightmat.update_usage(rowwise_usage=True) + + else: + weightmat = cast_if_needed(weightmat, activation_dtype) # Cast for AMP + # ------------------------------------------------------ + # Weight tensor is ready for GEMM... + # ------------------------------------------------------ + + # Cast bias to expected dtype + bias_dtype = activation_dtype + if needs_quantized_gemm(ln_out_total) and activation_dtype == torch.float32: + # cuBLAS does not support FP8 GEMM with FP32 bias, so we cast to BF16 + bias_dtype = torch.bfloat16 + bias_cast = cast_if_needed(bias, bias_dtype) if bias is not None else bias + + # Calibrate quantizers if needed + if not fp8 and args.fp8_calibration: + if input_quantizer is not None: + input_quantizer.calibrate(ln_out_total) + if weight_quantizer is not None: + weight_quantizer.calibrate(weight) - # Make sure input dimensions are compatible - out_features, in_features = weight.shape - inp_shape = inp.shape - inp_requires_grad = inp.requires_grad - assert inp_shape[-1] == in_features, "GEMM not possible" - inp = inp.view((-1, in_features)) - inputmat = inp - if fp8: - assert_dim_for_fp8_exec(inputmat, weight) + # Choose whether to use GEMM kernel with split accumulator + use_split_accumulator = _2X_ACC_FPROP + if fp8: + recipe = FP8GlobalStateManager.get_fp8_recipe() + if hasattr(recipe, "fp8_gemm_fprop"): + use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator + + # Configure output quantizer + if output_quantizer is not None: + output_quantizer.set_usage(rowwise=True, columnwise=False) + + # Output buffer for Userbuffers reduce-scatter + reduce_scatter_out = None + if ub_overlap_rs_fprop: + out_shape = list(inp_shape) + out_shape[0] //= tp_world_size + out_shape[-1] = out_features + reduce_scatter_out = torch.empty(out_shape, dtype=activation_dtype, device=inp.device) + + # ------------------------------------------------------ + # Forward GEMM + # Note: y = x * w^T + # ------------------------------------------------------ + nvtx_range_push(f"{nvtx_label}.gemm") + gemm_out, *_, reduce_scatter_out = general_gemm( + weightmat, + ln_out_total, + quantization_params=output_quantizer, + out_dtype=activation_dtype, + bias=bias_cast, + use_split_accumulator=use_split_accumulator, + ub=ub_obj, + ub_type=ub_type, + extra_output=reduce_scatter_out, + ) + nvtx_range_pop(f"{nvtx_label}.gemm") + # ------------------------------------------------------ + # Finished forward GEMM... + # ------------------------------------------------------ + + # Deallocate GEMM input tensor if no longer needed + if not weight_requires_grad and not return_layernorm_output: + clear_tensor_data(ln_out, ln_out_total) + ln_out = ln_out_total = None + elif with_input_all_gather and not return_layernorm_output_gathered: + # ln_out_total aliases ln_out for the cuBLASMp backend; skip the + # deallocation to avoid corrupting the backward-saved tensor. + if ln_out_total is not ln_out: + clear_tensor_data(ln_out_total) + ln_out_total = None - # Cast for native AMP - nvtx_range_push(f"{nvtx_label}.norm_input_cast") - inputmat = cast_if_needed(inputmat, activation_dtype) - ln_weight = cast_if_needed(ln_weight, activation_dtype) - if ln_bias is not None: - ln_bias = cast_if_needed(ln_bias, activation_dtype) - nvtx_range_pop(f"{nvtx_label}.norm_input_cast") + # ------------------------------------------------------ + # Prepare output tensor + # Note: Perform tensor-parallel communication + # ------------------------------------------------------ + out = None + if ub_overlap_rs_fprop: + # cuBLASMp writes the reduce-scattered output directly into the + # GEMM output tensor; Userbuffers writes it into the extra-output buffer. + out = gemm_out if ub_obj is not None and ub_obj.with_cublasmp() else reduce_scatter_out + elif parallel_mode == "row" and tp_size > 1: + nvtx_range_push(f"{nvtx_label}.row_parallel_comm") + out = gemm_out + if sequence_parallel: + out, _ = reduce_scatter_along_first_dim(out, tp_group) + elif args.tensor_parallel: + if args.symmetric_ar_type is not None: + out, _ = symmetric_all_reduce(out, tp_group, all_reduce_type=args.symmetric_ar_type) + else: + out, _ = allreduce(out, tp_group) + nvtx_range_pop(f"{nvtx_label}.row_parallel_comm") + else: + out = gemm_out + out = out.view(-1, *inp_shape[1:-1], out_features) + # ------------------------------------------------------ + # Output tensor is ready to return... + # ------------------------------------------------------ + + # Prepare backward state + tensors_to_save_from_forward = None + ctx_attrs = None + + if is_grad_enabled: + ln_out_to_save = ln_out + if backward_override == "high_precision": + ln_out_to_save = ln_out_hp + ln_out_needs_gather = ( + weight_requires_grad and parallel_mode == "column" and sequence_parallel + ) - if is_cpu_offload_enabled(): - start_offload(inputmat) + # Input with column-wise usage is needed for wgrad GEMM. + if backward_needs_input and backward_override is None: + if isinstance(ln_out, QuantizedTensorStorage): + # For sequence parallel in vanilla FP8, rowwise data is + # to gather the input. For MXFP8, columnwise only data + # can be allgathered. + if ( + isinstance(ln_out, (MXFP8TensorStorage, Float8BlockwiseQTensorStorage)) + or not ln_out_needs_gather + ): + ln_out.update_usage(rowwise_usage=False) + + if cpu_offloading: + mark_activation_offload(inputmat, mu, rsigma, ln_out_to_save) + + # Scatter intermediate/activation tensors saved for the backward pass + # NOTE: weight_fp8 = weight when fp8 == False and torch.disttributed.FSDP already + # shards/unshards the base weights so we don't do it ourselves + nvtx_range_push(f"{nvtx_label}.fsdp_scatter") + fsdp_shapes = _fsdp_scatter_tensors( + fsdp_group, + mu, + rsigma, + weightmat if fp8 and not is_weight_param_quantized else None, + ln_out_to_save if weight_requires_grad else None, + ) + nvtx_range_pop(f"{nvtx_label}.fsdp_scatter") - tp_world_size = get_distributed_world_size(tp_group) + if cpu_offloading: + mark_not_offload( + weightmat, + weight, + bias_cast, + ln_weight_cast, + ln_bias, + ) + + # FSDP2: Don't save FP8 workspace for non-quantized weights. + # Backward will re-quantize from FSDP2 all-gathered weight. + # (Issue #2681) + 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 + + # Dedup save slots that alias forward inputs or other outputs; + # ``_layernorm_linear_setup_ctx`` rebuilds the refs. + if wt_save is None: + wt_alias = None + elif wt_save is weight: + wt_alias = "weight" + elif new_weight_workspace is not None and wt_save is new_weight_workspace: + wt_alias = "new_weight_workspace" + elif args.weight_workspace is not None and wt_save is args.weight_workspace: + wt_alias = "weight_workspace" + else: + wt_alias = None + saved_tensor_aliases = ( + "inp" if inputmat is inp else None, + wt_alias, + "weight", # ``saved_weight`` slot is always the weight parameter + "bias" if bias_cast is not None and bias_cast is bias else None, + "ln_weight" if ln_weight_cast is ln_weight else None, + ( + "ln_out" + if return_layernorm_output + and ln_out_to_save is not None + and ln_out_to_save is ln_out_return + else None + ), + None, + None, + ) + tensors_to_save_from_forward = ( + None if saved_tensor_aliases[0] is not None else inputmat, + None if saved_tensor_aliases[1] is not None else wt_save, + None, + None if saved_tensor_aliases[3] is not None else bias_cast, + None if saved_tensor_aliases[4] is not None else ln_weight_cast, + None if saved_tensor_aliases[5] is not None else ln_out_to_save, + mu, + rsigma, + ) - weight_requires_grad = weight.requires_grad - backward_needs_input = is_grad_enabled and weight_requires_grad + ctx_attrs = { + "fsdp_shapes": fsdp_shapes, + "saved_tensor_aliases": saved_tensor_aliases, + "is_weight_param_quantized": is_weight_param_quantized, + "ln_out_needs_gather": ln_out_needs_gather, + } + + ln_out_for_return = None + if return_layernorm_output: + if return_layernorm_output_gathered: + shape = list(inp_shape) + shape[0] *= tp_size if with_input_all_gather else 1 + ln_out_for_return = ln_out_return.view(shape) + else: + ln_out_for_return = ln_out_return.view(inp_shape) + return out, ln_out_for_return, new_weight_workspace, tensors_to_save_from_forward, ctx_attrs + + +def _layernorm_linear_setup_ctx( + bwd_args: LayerNormLinearBwdArgs, + fwd_args: LayerNormLinearFwdArgs, + fwd_outputs: Tuple[Any, ...], + ctx_attrs: Dict, + tensors_to_save_from_forward: Tuple[Any, ...], +) -> Tuple[Any, ...]: + """Populate ``bwd_args`` from forward state. + + Returns the tensors that should be passed through ``prepare_for_saving`` + by the caller. ``fwd_outputs`` is ``(out, ln_out_return, + new_weight_workspace)``; the last two rebuild the deduped save slots. + """ + inp = fwd_args.inp + weight = fwd_args.weight + bias = fwd_args.bias + ln_weight = fwd_args.ln_weight + + backward_override = fwd_args.backward_override + fp8 = fwd_args.fp8 + debug = fwd_args.debug + fuse_wgrad_accumulation = fwd_args.fuse_wgrad_accumulation + is_weight_param_quantized = ctx_attrs["is_weight_param_quantized"] + + # Quantizers + bwd_args.input_quantizer = fwd_args.input_quantizer + bwd_args.weight_quantizer = ( + weight._quantizer + if (is_weight_param_quantized and not debug and isinstance(weight, QuantizedTensorStorage)) + else fwd_args.weight_quantizer + ) + bwd_args.grad_input_quantizer = fwd_args.grad_input_quantizer + bwd_args.grad_weight_quantizer = fwd_args.grad_weight_quantizer + bwd_args.grad_output_quantizer = fwd_args.grad_output_quantizer + + # Differentiability summary + bwd_args.use_bias = bias is not None + bwd_args.requires_dgrad = fwd_args.input_requires_grad + bwd_args.requires_wgrad = fwd_args.weight_requires_grad + bwd_args.ln_out_needs_gather = ctx_attrs["ln_out_needs_gather"] + bwd_args.inp_shape = inp.shape + + # Normalization + bwd_args.normalization = fwd_args.normalization + bwd_args.zero_centered_gamma = fwd_args.zero_centered_gamma + bwd_args.bwd_ln_sm_margin = fwd_args.bwd_ln_sm_margin + bwd_args.return_layernorm_output = fwd_args.return_layernorm_output + bwd_args.return_layernorm_output_gathered = fwd_args.return_layernorm_output_gathered + + # Numerical / dtype config + bwd_args.activation_dtype = fwd_args.activation_dtype + bwd_args.fp8 = fp8 + bwd_args.dgrad_use_split_accumulator = fwd_args.dgrad_use_split_accumulator + bwd_args.wgrad_use_split_accumulator = fwd_args.wgrad_use_split_accumulator + bwd_args.backward_override = backward_override + bwd_args.is_weight_param_quantized = is_weight_param_quantized + bwd_args.debug = debug + + # Tensor / sequence parallelism + bwd_args.parallel_mode = fwd_args.parallel_mode + bwd_args.tp_group = fwd_args.tp_group + bwd_args.tp_size = fwd_args.tp_size + bwd_args.tensor_parallel = fwd_args.tensor_parallel + bwd_args.sequence_parallel = fwd_args.sequence_parallel + + # Userbuffers + bwd_args.ub_name = fwd_args.ub_name + bwd_args.ub_overlap_ag = fwd_args.ub_overlap_ag_dgrad + bwd_args.ub_overlap_rs_dgrad = fwd_args.ub_overlap_rs_dgrad + bwd_args.ub_bulk_dgrad = fwd_args.ub_bulk_dgrad + bwd_args.ub_bulk_wgrad = fwd_args.ub_bulk_wgrad + + # FSDP + bwd_args.fsdp_group = fwd_args.fsdp_group + bwd_args.fsdp_shapes = ctx_attrs["fsdp_shapes"] + bwd_args.is_fsdp2 = fwd_args.is_fsdp2 + + # Weight-grad scheduling / accumulation + bwd_args.is_first_microbatch = fwd_args.is_first_microbatch + bwd_args.fuse_wgrad_accumulation = fuse_wgrad_accumulation + bwd_args.wgrad_store = fwd_args.wgrad_store + if fuse_wgrad_accumulation and fwd_args.weight_requires_grad: + # Keep weakref to weight to preserve attributes like main_grad + # when we need to modify the weight python object + bwd_args.origin_weight_ref = weakref.ref(weight) + bwd_args.origin_weight_overwrites_main_grad = getattr(weight, "overwrite_main_grad", False) + # MCore FSDP creates main_grad lazily before backward, so don't touch it here + 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 + + # Misc + bwd_args.cpu_offloading = fwd_args.cpu_offloading + + if backward_override is not None: + bwd_args.fp8 = False + bwd_args.debug = False + bwd_args.ub_overlap_ag = False + bwd_args.ub_overlap_rs_dgrad = False + bwd_args.ub_bulk_dgrad = False + bwd_args.ub_bulk_wgrad = False + bwd_args.grad_input_quantizer = None + bwd_args.grad_weight_quantizer = None + bwd_args.grad_output_quantizer = None + + ( + saved_inputmat, + wt_save, + saved_weight, + saved_bias, + saved_ln_weight, + saved_ln_out, + mu, + rsigma, + ) = tensors_to_save_from_forward + ( + inputmat_alias, + wt_save_alias, + saved_weight_alias, + bias_alias, + ln_weight_alias, + ln_out_alias, + _, + _, + ) = ctx_attrs["saved_tensor_aliases"] + in_features = inp.shape[-1] + if inputmat_alias == "inp": + saved_inputmat = inp.view((-1, in_features)) + if wt_save_alias == "weight": + wt_save = weight + elif wt_save_alias == "new_weight_workspace": + wt_save = fwd_outputs[2] + elif wt_save_alias == "weight_workspace": + wt_save = fwd_args.weight_workspace + if saved_weight_alias == "weight": + saved_weight = weight + if bias_alias == "bias": + saved_bias = bias + if ln_weight_alias == "ln_weight": + saved_ln_weight = ln_weight + if ln_out_alias == "ln_out": + saved_ln_out = fwd_outputs[1].view((-1, in_features)) + if fwd_args.cpu_offloading: + # Rebuilt views don't carry the offload marks set on the forward tensors + mark_activation_offload( + saved_inputmat if inputmat_alias == "inp" else None, + saved_ln_out if ln_out_alias == "ln_out" else None, + ) + return ( + saved_inputmat, + wt_save, + saved_weight, + saved_bias, + saved_ln_weight, + saved_ln_out, + mu, + rsigma, + ) + + +def _layernorm_linear_backward_impl( + args: LayerNormLinearBwdArgs, +) -> Tuple[Union[torch.Tensor, None], ...]: + """Backward implementation for the layernorm-linear layer. + + Caller must have populated ``args.grad_output`` / ``args.grad_ln_out`` and + run ``args.setup_saved_tensors(ctx)`` before invocation. Returns + ``(dgrad, dgamma, dbeta, wgrad, grad_bias)``. + """ + grad_output = args.grad_output + assert grad_output is not None + + # NVTX label for profiling + nvtx_label = "transformer_engine._LayerNormLinear.backward" + if args.ub_name is not None: + nvtx_label = f"{nvtx_label}.{args.ub_name}" + + with get_nvtx_range_context("_LayerNormLinear_backward"): + inputmat = args.inputmat + weight = args.weight_fp8 + saved_weight = args.saved_weight + bias = args.bias + ln_weight = args.ln_weight + ln_out = args.ln_out + mu = args.mu + rsigma = args.rsigma + + 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. + origin_weight = None + origin_weight_overwrites_main_grad = args.origin_weight_overwrites_main_grad + main_grad = None + if args.fuse_wgrad_accumulation and args.requires_wgrad: + origin_weight_ref = args.origin_weight_ref + args.origin_weight_ref = None + origin_weight = origin_weight_ref() if origin_weight_ref is not None else None + assert ( + origin_weight is not None + ), "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 = args.main_grad_func() if weight is not None else None + if main_grad is not None and not is_dist_weight: + origin_weight.main_grad = main_grad + + # Gather intermediate/activation tensors if needed + # NOTE: weight_fp8 = weight when fp8 == False and torch.disttributed.FSDP already + # shards/unshards the base weights so we don't do it ourselves + nvtx_range_push(f"{nvtx_label}.fsdp_gather") + _fsdp_gather_tensors( + args.fsdp_group, + args.fsdp_shapes, + mu, + rsigma, + weight if args.fp8 and not args.is_weight_param_quantized else None, + ln_out, + ) + nvtx_range_pop(f"{nvtx_label}.fsdp_gather") # Configure Userbuffers communication (comm+GEMM overlap) - if debug: # turn off userbuffers in debug mode - ub_overlap_ag_fprop = False - ub_overlap_rs_fprop = False - ub_overlap_ag_dgrad = False - ub_overlap_rs_dgrad = False - ub_bulk_wgrad = False - ub_bulk_dgrad = False - ub_obj = None - ub_type = None - ub_overlap_ag_fprop = ( - ub_overlap_ag_fprop and is_grad_enabled and not return_layernorm_output - ) - if ub_overlap_rs_fprop: - ub_obj = get_ub(ub_name + "_fprop", fp8) - ub_type = tex.CommOverlapType.RS - elif ub_overlap_ag_fprop: - ub_obj = get_ub(ub_name + "_fprop", fp8) - ub_type = tex.CommOverlapType.AG - - # Configure quantizer for norm output - if fp8: - if input_quantizer is None: - raise ValueError("Missing quantizer for input tensor") - input_quantizer.set_usage( - rowwise=True, - columnwise=backward_needs_input and backward_override is None, - ) - if with_input_all_gather and input_quantizer.supports_only_rowwise_all_gather(): - # All-gather is not supported with FP8 column-wise data - input_quantizer.set_usage(columnwise=False) - # Amax reduction group for the input quantizer (column-parallel sequence parallel) + args.ub_obj_gradout = None + ub_obj_dgrad = None + ub_obj_wgrad = None + ub_type_dgrad = None + ub_type_wgrad = None + dgrad_shape = [reduce(multiply_op, args.inp_shape[:-1]), args.inp_shape[-1]] + if args.ub_overlap_ag: + # Overlap grad_output all-gather with dgrad compute + args.ub_obj_gradout = get_ub(args.ub_name + "_dgrad", args.fp8) + ub_obj_dgrad = args.ub_obj_gradout + ub_type_dgrad = tex.CommOverlapType.AG + elif args.ub_overlap_rs_dgrad: + # Overlap dgrad reduce-scatter with dgrad compute + args.ub_obj_gradout = get_ub(args.ub_name + "_dgrad", args.fp8) + ub_obj_dgrad = args.ub_obj_gradout + ub_type_dgrad = tex.CommOverlapType.RS + else: + if args.ub_bulk_dgrad: + # Overlap inputmat all-gather with dgrad compute + args.ub_obj_gradout = get_ub(args.ub_name + "_dgrad", args.fp8) + ub_obj_dgrad = args.ub_obj_gradout + ub_type_dgrad = tex.CommOverlapType.AG + if args.ub_bulk_wgrad: + # Overlap dgrad reduce-scatter with wgrad compute + ub_obj_wgrad = get_ub(args.ub_name + "_wgrad", args.fp8) + ub_type_wgrad = tex.CommOverlapType.RS + + # -------------------------------------------------- + # Prepare grad output tensor + # Note: Cast to expected dtype and perform tensor-parallel communication + # -------------------------------------------------- + + # Configure quantizer for grad output tensor + # Note: dgrad GEMM requires row-wise usage, wgrad GEMM + # requires column-wise usage + if args.grad_output_quantizer is not None: + quantizer = args.grad_output_quantizer + quantizer.set_usage(rowwise=True, columnwise=True) + if args.ub_overlap_ag: + # Userbuffers only supports communication for one + # tensor usage at a time. Configure quantizer with + # usage for only dgrad GEMM. + quantizer.set_usage(columnwise=False) + # Amax reduction group for grad output (row-parallel sequence parallel) set_quantizer_amax_reduction_group( - input_quantizer, - tp_group if (sequence_parallel and parallel_mode == "column") else None, + quantizer, + ( + args.tp_group + if (args.sequence_parallel and args.parallel_mode == "row") + else None + ), ) - # 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 - and not return_layernorm_output - 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 - ) - - # Apply normalization - nvtx_range_push(f"{nvtx_label}.norm") - ln_out, mu, rsigma = apply_normalization( - inputmat, - None, # ln_out - ln_weight, - ln_bias, - eps, - input_quantizer if with_quantized_norm else None, - inputmat.dtype, - normalization, - fwd_ln_sm_margin, - zero_centered_gamma, + # Prepare grad output tensor + # Note: Cast to expected dtype and perform tensor-parallel communication + nvtx_range_push(f"{nvtx_label}.grad_output_preprocess") + ( + grad_output, + grad_bias, + ) = TransformerEngineBaseModule.grad_output_preprocess( + args, + grad_output, + args.parallel_mode == "row", + args.grad_output_quantizer, ) - nvtx_range_pop(f"{nvtx_label}.norm") + nvtx_range_pop(f"{nvtx_label}.grad_output_preprocess") - # Store unquantized layer norm output if we need to return it - ln_out_return = None - if return_layernorm_output or return_layernorm_output_gathered: - ln_out_return = ln_out - ln_out_hp = ln_out if backward_override == "high_precision" else None + # -------------------------------------------------- + # Grad output tensor is ready for computing grad input... + # -------------------------------------------------- - # ------------------------------------------------------ + # -------------------------------------------------- # Prepare GEMM input tensor - # Note: Cast to expected dtype and perform tensor-parallel communication - # ------------------------------------------------------ - nvtx_range_push(f"{nvtx_label}.gemm_input_cast_comm") + # Note: Input tensor is needed for wgrad GEMM. + # Tensor-parallel communication is overlapped with dgrad + # GEMM. + # -------------------------------------------------- ln_out_total = None - if with_input_all_gather: - if return_layernorm_output_gathered: - # Perform all-gather in high precision if gathered - # norm output will be returned - ln_out_total, _ = gather_along_first_dim(ln_out, tp_group) - ln_out_return = ln_out_total - if fp8 or debug: - ln_out = input_quantizer(ln_out) - input_quantizer.set_usage(rowwise=True, columnwise=False) - ln_out_total = input_quantizer(ln_out_total) + ln_out_total_work = None + if args.backward_override == "dequantized": + if isinstance(ln_out, QuantizedTensorStorage): + ln_out = ln_out.dequantize(dtype=args.activation_dtype) else: - quantizer = None - if fp8 or debug: - quantizer = input_quantizer - # custom recipe doesn't need to support quantized AG - if not with_quantized_norm and not custom: - ln_out = quantizer(ln_out) - quantizer.set_usage(rowwise=True, columnwise=False) - if ub_overlap_ag_fprop: # Initialize Userbuffers all-gather - ln_out_total, _ = fill_userbuffers_buffer_for_all_gather( - ub_obj, - ln_out, - quantizer, - tp_group, - ) - else: # Perform NCCL all-gather - ln_out_total, _ = gather_along_first_dim( - ln_out, - tp_group, - quantizer=quantizer, - ) - else: - if (fp8 or debug) and not with_quantized_norm: - ln_out = input_quantizer(ln_out) - ln_out_total = ln_out - nvtx_range_pop(f"{nvtx_label}.gemm_input_cast_comm") - # ------------------------------------------------------ - # GEMM input tensor is ready... - # ------------------------------------------------------ - - # ------------------------------------------------------ - # 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 - if fp8 or debug: - is_weight_param_quantized = isinstance(weight, QuantizedTensorStorage) - - # Configure quantizer - # If weight is already quantized, weight._quantizer is its true quantizer. - # for debug mode we create quantizer every iteration, thus we need to set the quantizer states - if is_weight_param_quantized and not debug: - weight_quantizer = weight._quantizer - elif weight_quantizer is not None: - # FSDP2: Skip columnwise/transpose creation during forward - # to avoid accumulating caches across layers. Backward's - # FSDP2 all-gather will recreate them. (Issue #2681) - weight_quantizer.set_usage( - rowwise=True, - columnwise=is_grad_enabled and not is_fsdp2 and backward_override is None, + ln_out = cast_if_needed(ln_out, args.activation_dtype) + if args.ln_out_needs_gather: + quantizer = None + if args.input_quantizer is not None and args.fp8: + quantizer = args.input_quantizer + set_quantizer_usage_for_wgrad_all_gather(quantizer) + if args.ub_bulk_dgrad: + ln_out_total, _ = fill_userbuffers_buffer_for_all_gather( + ub_obj_dgrad, + ln_out, + quantizer, + args.tp_group, ) - - # Get quantized weight - update_ws = is_first_microbatch is None or is_first_microbatch - weightmat, new_weight_workspace = quantize_weight( - tensor=weight, - quantizer=weight_quantizer, - workspace=weight_workspace, - update_workspace=update_ws, - skip_update_flag=skip_fp8_weight_update, - fsdp_group=fsdp_group, - workspace_dtype=activation_dtype, - cache=cache_weight, - ) - - weightmat.update_usage(rowwise_usage=True) - + else: + nvtx_range_push(f"{nvtx_label}.column_parallel_comm_input") + ln_out_total, ln_out_total_work = gather_along_first_dim( + ln_out, + args.tp_group, + async_op=True, + quantizer=quantizer, + ) + nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_input") else: - weightmat = cast_if_needed(weightmat, activation_dtype) # Cast for AMP - # ------------------------------------------------------ - # Weight tensor is ready for GEMM... - # ------------------------------------------------------ - - # Cast bias to expected dtype - bias_dtype = activation_dtype - if needs_quantized_gemm(ln_out_total) and activation_dtype == torch.float32: - # cuBLAS does not support FP8 GEMM with FP32 bias, so we cast to BF16 - bias_dtype = torch.bfloat16 - bias = cast_if_needed(bias, bias_dtype) if bias is not None else bias - - # Calibrate quantizers if needed - if not fp8 and fp8_calibration: - if input_quantizer is not None: - input_quantizer.calibrate(ln_out_total) - if weight_quantizer is not None: - weight_quantizer.calibrate(weight) + ln_out_total = ln_out + # -------------------------------------------------- + # Input tensor is ready for computing grad weight... + # -------------------------------------------------- + + # -------------------------------------------------- + # Compute grad input tensor + # Note: Gradient w.r.t. GEMM input (i.e. norm output). + # -------------------------------------------------- + + # FSDP2: Re-create workspace from all-gathered weight when + # workspace was not saved. (Issue #2681) + # Use saved_weight (the original weight parameter) since + # origin_weight is only set when fuse_wgrad_accumulation=True. + if weight is None: + if isinstance(saved_weight, QuantizedTensorStorage): + # saved weight is already set to right usages by + # fsdp2 quantized-tensor hooks when workspace was not saved. + weight = saved_weight + elif args.weight_quantizer is not None: + args.weight_quantizer.set_usage(rowwise=True, columnwise=True) + weight = args.weight_quantizer(saved_weight) + + # Make sure required data is available + if isinstance(grad_output, QuantizedTensorStorage): + grad_output.update_usage(rowwise_usage=True) + if ( + args.fp8 + and args.weight_quantizer is not None + and isinstance(weight, QuantizedTensorStorage) + ): + weight.update_usage(columnwise_usage=True) # Choose whether to use GEMM kernel with split accumulator - use_split_accumulator = _2X_ACC_FPROP - if fp8: - recipe = FP8GlobalStateManager.get_fp8_recipe() - if hasattr(recipe, "fp8_gemm_fprop"): - use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator + use_split_accumulator = args.dgrad_use_split_accumulator - # Configure output quantizer - if output_quantizer is not None: - output_quantizer.set_usage(rowwise=True, columnwise=False) + # Update grad input quantizer + if args.grad_input_quantizer is not None: + args.grad_input_quantizer.set_usage(rowwise=True, columnwise=False) - # Output buffer for Userbuffers reduce-scatter + # Output buffers for Userbuffers reduce-scatter + gemm_out = None reduce_scatter_out = None - if ub_overlap_rs_fprop: - out_shape = list(inp_shape) - out_shape[0] //= tp_world_size - out_shape[-1] = out_features - reduce_scatter_out = torch.empty(out_shape, dtype=activation_dtype, device=inp.device) - - # ------------------------------------------------------ - # Forward GEMM - # Note: y = x * w^T - # ------------------------------------------------------ - nvtx_range_push(f"{nvtx_label}.gemm") + if args.ub_overlap_rs_dgrad: + reduce_scatter_out = torch.empty( + dgrad_shape, dtype=args.activation_dtype, device=args.grad_output.device + ) + elif args.ub_bulk_wgrad: + gemm_out = ub_obj_wgrad.get_buffer(local_chunk=False) + + # dgrad GEMM + # Note: dx = dy * w + nvtx_range_push(f"{nvtx_label}.dgrad_gemm") + weight_for_dgrad = weight + if args.backward_override == "dequantized": + if isinstance(weight_for_dgrad, QuantizedTensorStorage): + weight_for_dgrad = weight_for_dgrad.dequantize(dtype=args.activation_dtype) + else: + weight_for_dgrad = cast_if_needed(weight_for_dgrad, args.activation_dtype) + elif args.backward_override == "high_precision": + weight_for_dgrad = saved_weight + if isinstance(weight_for_dgrad, QuantizedTensorStorage): + weight_for_dgrad = weight_for_dgrad.dequantize(dtype=args.activation_dtype) gemm_out, *_, reduce_scatter_out = general_gemm( - weightmat, - ln_out_total, - quantization_params=output_quantizer, - out_dtype=activation_dtype, - bias=bias, + weight_for_dgrad, + grad_output, + layout="NN", + grad=True, + quantization_params=args.grad_input_quantizer, + out=gemm_out, + out_dtype=args.activation_dtype, use_split_accumulator=use_split_accumulator, - ub=ub_obj, - ub_type=ub_type, + ub=ub_obj_dgrad, + ub_type=ub_type_dgrad, extra_output=reduce_scatter_out, + bulk_overlap=args.ub_bulk_dgrad, ) - nvtx_range_pop(f"{nvtx_label}.gemm") - # ------------------------------------------------------ - # Finished forward GEMM... - # ------------------------------------------------------ - - # Deallocate GEMM input tensor if no longer needed - if not weight.requires_grad and not return_layernorm_output: - clear_tensor_data(ln_out, ln_out_total) - ln_out = ln_out_total = None - elif with_input_all_gather and not return_layernorm_output_gathered: - # ln_out_total aliases ln_out for the cuBLASMp backend; skip the - # deallocation to avoid corrupting the backward-saved tensor. - if ln_out_total is not ln_out: - clear_tensor_data(ln_out_total) - ln_out_total = None - - # ------------------------------------------------------ - # Prepare output tensor - # Note: Perform tensor-parallel communication - # ------------------------------------------------------ - out = None - if ub_overlap_rs_fprop: - # cuBLASMp writes the reduce-scattered output directly into the - # GEMM output tensor; Userbuffers writes it into the extra-output buffer. - out = gemm_out if ub_obj is not None and ub_obj.with_cublasmp() else reduce_scatter_out - elif parallel_mode == "row" and tp_size > 1: - nvtx_range_push(f"{nvtx_label}.row_parallel_comm") - out = gemm_out - if sequence_parallel: - out, _ = reduce_scatter_along_first_dim(out, tp_group) - elif tensor_parallel: - if symmetric_ar_type is not None: - out, _ = symmetric_all_reduce(out, tp_group, all_reduce_type=symmetric_ar_type) - else: - out, _ = allreduce(out, tp_group) - nvtx_range_pop(f"{nvtx_label}.row_parallel_comm") - else: - out = gemm_out - out = out.view(-1, *inp_shape[1:-1], out_features) - # ------------------------------------------------------ - # Output tensor is ready to return... - # ------------------------------------------------------ - - # ------------------------------------------------------ - # Cache state for backward pass - # ------------------------------------------------------ + nvtx_range_pop(f"{nvtx_label}.dgrad_gemm") - if is_grad_enabled: - ln_out_to_save = ln_out - if backward_override == "high_precision": - ln_out_to_save = ln_out_hp - ctx.weight_quantizer = weight_quantizer - ctx.ln_out_needs_gather = ( - weight.requires_grad and parallel_mode == "column" and sequence_parallel - ) + # FSDP2 only handles deallocation all-gathered weights that it allocates. + # Columnwise data is derived from rowwise data after allgather for fp8 + # and 2d block-scaled weights in TE managed memory. So we need to clear + # it here. + # (Issues #2681, #2717) + if args.is_fsdp2 and isinstance(weight, QuantizedTensorStorage): + clear_columnwise_cache(weight) - # Input with column-wise usage is needed for wgrad GEMM. - if backward_needs_input and backward_override is None: - if isinstance(ln_out, QuantizedTensorStorage): - # For sequence parallel in vanilla FP8, rowwise data is - # to gather the input. For MXFP8, columnwise only data - # can be allgathered. - if ( - isinstance(ln_out, (MXFP8TensorStorage, Float8BlockwiseQTensorStorage)) - or not ctx.ln_out_needs_gather - ): - ln_out.update_usage(rowwise_usage=False) - - if cpu_offloading: - mark_activation_offload(inputmat, mu, rsigma, ln_out_to_save) - - # Scatter intermediate/activation tensors saved for the backward pass - # NOTE: weight_fp8 = weight when ctx.fp8 == False and torch.disttributed.FSDP already - # shards/unshards the base weights so we don't do it ourselves - nvtx_range_push(f"{nvtx_label}.fsdp_scatter") - ctx.fsdp_group = fsdp_group - ctx.fsdp_shapes = _fsdp_scatter_tensors( - fsdp_group, - mu, - rsigma, - weightmat if fp8 and not is_weight_param_quantized else None, - ln_out_to_save if weight.requires_grad else None, + # Prepare grad input tensor + # Note: Perform tensor-parallel communication + dgrad = None + dgrad_work = None + if args.ub_overlap_rs_dgrad: + # cuBLASMp writes the reduce-scattered dgrad directly into the + # GEMM output tensor; Userbuffers uses the extra-output buffer. + dgrad = ( + gemm_out + if ub_obj_dgrad is not None and ub_obj_dgrad.with_cublasmp() + else reduce_scatter_out ) - nvtx_range_pop(f"{nvtx_label}.fsdp_scatter") - - if cpu_offloading: - mark_not_offload( - weightmat, - weight, - bias, - ln_weight, - ln_bias, + elif args.ub_bulk_wgrad: + dgrad = ub_obj_wgrad.get_buffer(local_chunk=True) + elif args.parallel_mode == "column" and args.tp_size > 1: + nvtx_range_push(f"{nvtx_label}.column_parallel_comm_dgrad") + dgrad = gemm_out + if args.sequence_parallel: + dgrad, dgrad_work = reduce_scatter_along_first_dim( + dgrad, + args.tp_group, + async_op=True, ) + else: + dgrad, dgrad_work = allreduce(dgrad, args.tp_group, async_op=True) + nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_dgrad") + else: + dgrad = gemm_out + + # -------------------------------------------------- + # Grad input tensor has been computed... + # -------------------------------------------------- + + # cuBLASMp's AG+GEMM consumes the gathered grad_output inline and + # does not preserve it for wgrad. Userbuffers leaves the gathered + # tensor in its persistent buffer; cuBLASMp does not, so we gather + # here. Route through the same FP8-aware all-gather as the + # non-overlap path in + # ``TransformerEngineBaseModule.grad_output_preprocess`` by passing + # the grad_output quantizer. Columnwise data needed for wgrad is + # produced by ``update_usage(columnwise_usage=True)`` further below. + if ( + args.requires_wgrad + and args.ub_overlap_ag + and args.ub_obj_gradout is not None + and args.ub_obj_gradout.with_cublasmp() + ): + if args.grad_output_quantizer is not None: + set_quantizer_usage_for_wgrad_all_gather(args.grad_output_quantizer) + grad_output, _ = gather_along_first_dim( + grad_output, + args.tp_group, + quantizer=args.grad_output_quantizer, + ) - # FSDP2: Don't save FP8 workspace for non-quantized weights. - # Backward will re-quantize from FSDP2 all-gathered weight. - # (Issue #2681) - 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 + # -------------------------------------------------- + # Compute grad weight + # -------------------------------------------------- - tensors_to_save, tensor_objects = prepare_for_saving( - inputmat, - wt_save, - origin_weight, - bias, - ln_weight, - ln_out_to_save, - mu, - rsigma, - ) - ctx.save_for_backward(*tensors_to_save) - ctx.tensor_objects = tensor_objects - ctx.requires_dgrad = inp_requires_grad - ctx.requires_wgrad = weight.requires_grad - ctx.is_weight_param_quantized = is_weight_param_quantized - ctx.is_fsdp2 = is_fsdp2 - if fuse_wgrad_accumulation and weight.requires_grad: - # Keep weakref to weight to preserve attributes like main_grad - # when we need to modify the weight python object - ctx.origin_weight_ref = weakref.ref(weight) - # Save overwrite_main_grad flag now while we have access to weight object - ctx.origin_weight_overwrites_main_grad = getattr( - weight, "overwrite_main_grad", False + wgrad = None + if args.requires_wgrad: + # Prepare grad output tensor + # Note: Synchronize tensor-parallel communication and + # make sure required data is available + if args.ub_overlap_ag and isinstance(args.grad_output_quantizer, MXFP8Quantizer): + # UB does not support pipelined overlapping grad output + # all-gather with wgrad GEMM. Also, we can't + # convert row-scaled MXFP8 to column-scaled, so we + # can't reuse the grad output that was gathered + # for the dgrad GEMM. We work around by explicitly + # overlapping the AG operation with the dgrad GEMM. + + # Get the communication stream from the dgrad GEMM to use for the AG + dgrad_send_stream, dgrad_recv_stream = ub_obj_dgrad.get_communication_stream() + + # This object is separate from the ub_obj_wgrad object which is passed to the GEMM + ub_obj_overlap_wgrad = get_ub(args.ub_name + "_wgrad", args.fp8) + + args.grad_output_quantizer.set_usage(rowwise=False, columnwise=True) + + # We use the send stream to copy into the userbuffers. + # This is the same stream that we will use to access the data in the AG, + # so we dont need to add any syncs yet. + with torch.cuda.stream(dgrad_send_stream): + grad_output, _ = fill_userbuffers_buffer_for_all_gather( + ub_obj_overlap_wgrad, + args.grad_output, + args.grad_output_quantizer, + args.tp_group, + ) + + # Allgather grad_outputs[0] using the dgrad streams so we can overlap with the fc2_dgrad gemm + tex.bulk_overlap_ag_with_external_gemm( + ub_obj_overlap_wgrad, dgrad_send_stream, dgrad_recv_stream ) - # This check is needed to ensure that main_grad is not created - # during the forward pass when using MCore FSDP as it creates - # the main_grad buffer lazily before backprop - 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 - ctx.grad_weight_quantizer = grad_weight_quantizer - ctx.grad_output_quantizer = grad_output_quantizer - ctx.input_quantizer = input_quantizer - ctx.owns_input = inputmat is not inp - ctx.weight = weight - ctx.activation_dtype = activation_dtype - ctx.fp8 = fp8 - ctx.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None - ctx.backward_override = backward_override - ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation - ctx.cpu_offloading = cpu_offloading - ctx.is_first_microbatch = is_first_microbatch - ctx.use_bias = bias is not None - ctx.sequence_parallel = sequence_parallel - ctx.tensor_parallel = tensor_parallel - ctx.inp_shape = inp_shape - ctx.parallel_mode = parallel_mode - ctx.tp_group = tp_group - ctx.tp_size = tp_size - ctx.return_layernorm_output = return_layernorm_output - ctx.return_layernorm_output_gathered = return_layernorm_output_gathered - ctx.bwd_ln_sm_margin = bwd_ln_sm_margin - ctx.zero_centered_gamma = zero_centered_gamma - ctx.ub_overlap_ag = ub_overlap_ag_dgrad - ctx.ub_overlap_rs_dgrad = ub_overlap_rs_dgrad - ctx.ub_bulk_wgrad = ub_bulk_wgrad - ctx.ub_bulk_dgrad = ub_bulk_dgrad - ctx.ub_name = ub_name - ctx.requires_dgrad = inp_requires_grad - ctx.normalization = normalization - ctx.reduce_and_update_bwd_fp8_tensors = False - if ctx.fp8 and requires_grad(inp, ln_weight, ln_bias, weight, bias): - qstate = FP8GlobalStateManager.quantization_state - _first_fp8_module = qstate.is_first_fp8_module - ctx.reduce_and_update_bwd_fp8_tensors = FP8GlobalStateManager.is_first_fp8_module() - if in_fp8_activation_recompute_phase(): - qstate.is_first_fp8_module = _first_fp8_module - ctx.wgrad_store = wgrad_store - ctx.debug = debug - - # backward overrides - if backward_override is not None: - ctx.fp8 = False - ctx.debug = False - ctx.ub_overlap_ag = False - ctx.ub_overlap_rs_dgrad = False - ctx.ub_bulk_dgrad = False - ctx.ub_bulk_wgrad = False - ctx.grad_input_quantizer = None - ctx.grad_weight_quantizer = None - ctx.grad_output_quantizer = None - ctx.reduce_and_update_bwd_fp8_tensors = False - - # ------------------------------------------------------ - # Cached state for backward pass is ready... - # ------------------------------------------------------ - - ln_out_for_return = None - if return_layernorm_output: - if return_layernorm_output_gathered: - shape = list(inp_shape) - shape[0] *= tp_size if with_input_all_gather else 1 - ln_out_for_return = ln_out_return.view(shape) - else: - ln_out_for_return = ln_out_return.view(inp_shape) - return out, ln_out_for_return, new_weight_workspace - @staticmethod - def backward( - ctx, *grad_outputs: Tuple[torch.Tensor, ...] - ) -> Tuple[Union[torch.Tensor, None], ...]: - # pylint: disable=missing-function-docstring + # Prepare input tensor + # Note: Synchronize tensor-parallel communication and + # make sure required data is available + if ln_out_total_work is not None: + ln_out_total_work.wait() + ln_out_total_work = None + if args.fp8 or args.debug: + if isinstance(ln_out_total, QuantizedTensorStorage): + ln_out_total.update_usage(columnwise_usage=True) + else: + args.input_quantizer.set_usage(rowwise=False, columnwise=True) + ln_out_total = args.input_quantizer(ln_out_total) - # NVTX label for profiling - nvtx_label = "transformer_engine._LayerNormLinear.backward" - if ctx.ub_name is not None: - nvtx_label = f"{nvtx_label}.{ctx.ub_name}" + if args.fp8 or args.debug: + if isinstance(grad_output, QuantizedTensorStorage): + grad_output.update_usage(columnwise_usage=True) + else: + args.grad_output_quantizer.set_usage(rowwise=False, columnwise=True) + grad_output = args.grad_output_quantizer(grad_output) - with get_nvtx_range_context("_LayerNormLinear_backward"): - ( # pylint: disable=unbalanced-tuple-unpacking - inputmat, - weight, - saved_weight, - bias, - ln_weight, - ln_out, - mu, - rsigma, - ) = restore_from_func_ctx(ctx) + # Figure out whether to use split accumulator + use_split_accumulator = args.wgrad_use_split_accumulator - is_dist_weight = is_distributed_weight(saved_weight) + # Figure out whether to output wgrad GEMM directly into main grad 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. - origin_weight = None - origin_weight_overwrites_main_grad = getattr( - ctx, "origin_weight_overwrites_main_grad", False - ) - main_grad = None - if ctx.fuse_wgrad_accumulation and ctx.requires_wgrad: - origin_weight_ref = ctx.origin_weight_ref - ctx.origin_weight_ref = None - origin_weight = origin_weight_ref() if origin_weight_ref is not None else None - assert ( - origin_weight is not None - ), "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 and not is_dist_weight: - origin_weight.main_grad = main_grad - - # Gather intermediate/activation tensors if needed - # NOTE: weight_fp8 = weight when ctx.fp8 == False and torch.disttributed.FSDP already - # shards/unshards the base weights so we don't do it ourselves - nvtx_range_push(f"{nvtx_label}.fsdp_gather") - _fsdp_gather_tensors( - ctx.fsdp_group, - ctx.fsdp_shapes, - mu, - rsigma, - weight if ctx.fp8 and not ctx.is_weight_param_quantized else None, - ln_out, - ) - nvtx_range_pop(f"{nvtx_label}.fsdp_gather") - - # Configure Userbuffers communication (comm+GEMM overlap) - ctx.ub_obj_gradout = None - ub_obj_dgrad = None - ub_obj_wgrad = None - ub_type_dgrad = None - ub_type_wgrad = None - dgrad_shape = [reduce(multiply_op, ctx.inp_shape[:-1]), ctx.inp_shape[-1]] - if ctx.ub_overlap_ag: - # Overlap grad_output all-gather with dgrad compute - ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad", ctx.fp8) - ub_obj_dgrad = ctx.ub_obj_gradout - ub_type_dgrad = tex.CommOverlapType.AG - elif ctx.ub_overlap_rs_dgrad: - # Overlap dgrad reduce-scatter with dgrad compute - ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad", ctx.fp8) - ub_obj_dgrad = ctx.ub_obj_gradout - ub_type_dgrad = tex.CommOverlapType.RS - else: - if ctx.ub_bulk_dgrad: - # Overlap inputmat all-gather with dgrad compute - ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad", ctx.fp8) - ub_obj_dgrad = ctx.ub_obj_gradout - ub_type_dgrad = tex.CommOverlapType.AG - if ctx.ub_bulk_wgrad: - # Overlap dgrad reduce-scatter with wgrad compute - ub_obj_wgrad = get_ub(ctx.ub_name + "_wgrad", ctx.fp8) - ub_type_wgrad = tex.CommOverlapType.RS - - # -------------------------------------------------- - # Prepare grad output tensor - # Note: Cast to expected dtype and perform tensor-parallel communication - # -------------------------------------------------- - - # Configure quantizer for grad output tensor - # Note: dgrad GEMM requires row-wise usage, wgrad GEMM - # requires column-wise usage - if ctx.grad_output_quantizer is not None: - quantizer = ctx.grad_output_quantizer - quantizer.set_usage(rowwise=True, columnwise=True) - if ctx.ub_overlap_ag: - # Userbuffers only supports communication for one - # tensor usage at a time. Configure quantizer with - # usage for only dgrad GEMM. - quantizer.set_usage(columnwise=False) - # Amax reduction group for grad output (row-parallel sequence parallel) - set_quantizer_amax_reduction_group( - quantizer, - ( - ctx.tp_group - if (ctx.sequence_parallel and ctx.parallel_mode == "row") - else None - ), + # Distributed weight (e.g. GTP): accumulation happens downstream in finalize. + accumulate_wgrad_into_param_main_grad = False + elif args.is_first_microbatch is not None: + accumulate_wgrad_into_param_main_grad = ( + args.fuse_wgrad_accumulation and not args.is_first_microbatch ) - - # Prepare grad output tensor - # Note: Cast to expected dtype and perform tensor-parallel communication - nvtx_range_push(f"{nvtx_label}.grad_output_preprocess") - ( - grad_output, - grad_bias, - ) = TransformerEngineBaseModule.grad_output_preprocess( - ctx, - grad_outputs[0], - ctx.parallel_mode == "row", - ctx.grad_output_quantizer, - ) - nvtx_range_pop(f"{nvtx_label}.grad_output_preprocess") - - # -------------------------------------------------- - # Grad output tensor is ready for computing grad input... - # -------------------------------------------------- - - # -------------------------------------------------- - # Prepare GEMM input tensor - # Note: Input tensor is needed for wgrad GEMM. - # Tensor-parallel communication is overlapped with dgrad - # GEMM. - # -------------------------------------------------- - ln_out_total = None - ln_out_total_work = None - if ctx.backward_override == "dequantized": - if isinstance(ln_out, QuantizedTensorStorage): - ln_out = ln_out.dequantize(dtype=ctx.activation_dtype) - else: - ln_out = cast_if_needed(ln_out, ctx.activation_dtype) - if ctx.ln_out_needs_gather: - quantizer = None - if ctx.input_quantizer is not None and ctx.fp8: - quantizer = ctx.input_quantizer - 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, - ln_out, - quantizer, - ctx.tp_group, - ) - else: - nvtx_range_push(f"{nvtx_label}.column_parallel_comm_input") - ln_out_total, ln_out_total_work = gather_along_first_dim( - ln_out, - ctx.tp_group, - async_op=True, - quantizer=quantizer, - ) - nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_input") else: - ln_out_total = ln_out - # -------------------------------------------------- - # Input tensor is ready for computing grad weight... - # -------------------------------------------------- - - # -------------------------------------------------- - # Compute grad input tensor - # Note: Gradient w.r.t. GEMM input (i.e. norm output). - # -------------------------------------------------- - - # FSDP2: Re-create workspace from all-gathered weight when - # workspace was not saved. (Issue #2681) - # Use saved_weight (the original weight parameter) since - # origin_weight is only set when fuse_wgrad_accumulation=True. - if weight is None: - if isinstance(saved_weight, QuantizedTensorStorage): - # saved weight is already set to right usages by - # fsdp2 quantized-tensor hooks when workspace was not saved. - weight = saved_weight - elif ctx.weight_quantizer is not None: - ctx.weight_quantizer.set_usage(rowwise=True, columnwise=True) - weight = ctx.weight_quantizer(saved_weight) - - # Make sure required data is available - if isinstance(grad_output, QuantizedTensorStorage): - grad_output.update_usage(rowwise_usage=True) - if ( - ctx.fp8 - and ctx.weight_quantizer is not None - and isinstance(weight, QuantizedTensorStorage) - ): - weight.update_usage(columnwise_usage=True) - - # Choose whether to use GEMM kernel with split accumulator - use_split_accumulator = _2X_ACC_DGRAD - if ctx.fp8: - recipe = ctx.fp8_recipe - if hasattr(recipe, "fp8_gemm_dgrad"): - use_split_accumulator = recipe.fp8_gemm_dgrad.use_split_accumulator - - # Update grad input quantizer - if ctx.grad_input_quantizer is not None: - ctx.grad_input_quantizer.set_usage(rowwise=True, columnwise=False) - - # Output buffers for Userbuffers reduce-scatter - gemm_out = None + accumulate_wgrad_into_param_main_grad = args.fuse_wgrad_accumulation + + # Output buffer for overlapping FP8 grad input + # reduce-scatter with wgrad GEMM reduce_scatter_out = None - if ctx.ub_overlap_rs_dgrad: + if args.ub_bulk_wgrad and ub_obj_wgrad.is_fp8_ubuf(): reduce_scatter_out = torch.empty( - dgrad_shape, dtype=ctx.activation_dtype, device=grad_outputs[0].device - ) - elif ctx.ub_bulk_wgrad: - gemm_out = ub_obj_wgrad.get_buffer(local_chunk=False) - - # dgrad GEMM - # Note: dx = dy * w - nvtx_range_push(f"{nvtx_label}.dgrad_gemm") - weight_for_dgrad = weight - if ctx.backward_override == "dequantized": - if isinstance(weight_for_dgrad, QuantizedTensorStorage): - weight_for_dgrad = weight_for_dgrad.dequantize(dtype=ctx.activation_dtype) - else: - weight_for_dgrad = cast_if_needed(weight_for_dgrad, ctx.activation_dtype) - elif ctx.backward_override == "high_precision": - weight_for_dgrad = saved_weight - if isinstance(weight_for_dgrad, QuantizedTensorStorage): - weight_for_dgrad = weight_for_dgrad.dequantize(dtype=ctx.activation_dtype) - gemm_out, *_, reduce_scatter_out = general_gemm( - weight_for_dgrad, - grad_output, - layout="NN", - grad=True, - quantization_params=ctx.grad_input_quantizer, - out=gemm_out, - out_dtype=ctx.activation_dtype, - use_split_accumulator=use_split_accumulator, - ub=ub_obj_dgrad, - ub_type=ub_type_dgrad, - extra_output=reduce_scatter_out, - bulk_overlap=ctx.ub_bulk_dgrad, - ) - nvtx_range_pop(f"{nvtx_label}.dgrad_gemm") - - # FSDP2 only handles deallocation all-gathered weights that it allocates. - # Columnwise data is derived from rowwise data after allgather for fp8 - # and 2d block-scaled weights in TE managed memory. So we need to clear - # it here. - # (Issues #2681, #2717) - if getattr(ctx, "is_fsdp2", False) and isinstance(weight, QuantizedTensorStorage): - clear_columnwise_cache(weight) - - # Prepare grad input tensor - # Note: Perform tensor-parallel communication - dgrad = None - dgrad_work = None - if ctx.ub_overlap_rs_dgrad: - # cuBLASMp writes the reduce-scattered dgrad directly into the - # GEMM output tensor; Userbuffers uses the extra-output buffer. - dgrad = ( - gemm_out - if ub_obj_dgrad is not None and ub_obj_dgrad.with_cublasmp() - else reduce_scatter_out + dgrad_shape, dtype=args.activation_dtype, device=args.grad_output.device ) - elif ctx.ub_bulk_wgrad: - dgrad = ub_obj_wgrad.get_buffer(local_chunk=True) - elif ctx.parallel_mode == "column" and ctx.tp_size > 1: - nvtx_range_push(f"{nvtx_label}.column_parallel_comm_dgrad") - dgrad = gemm_out - if ctx.sequence_parallel: - dgrad, dgrad_work = reduce_scatter_along_first_dim( - dgrad, - ctx.tp_group, - async_op=True, + + # Arguments to include in wgrad GEMM closure + wgrad_gemm_kwargs = { + "out_dtype": ( + main_grad.dtype if args.fuse_wgrad_accumulation else args.activation_dtype + ), + "quantization_params": args.grad_weight_quantizer, + "accumulate": ( + accumulate_wgrad_into_param_main_grad + if not origin_weight_overwrites_main_grad + else False + ), + "layout": "NT", + "out": main_grad if args.fuse_wgrad_accumulation else None, + "bias": (bias if (grad_bias is None and not args.fp8) else None), + "use_split_accumulator": use_split_accumulator, + "grad": True, + "ub": ub_obj_wgrad, + "ub_type": ub_type_wgrad, + "extra_output": reduce_scatter_out, + "bulk_overlap": args.ub_bulk_wgrad, + } + + def wgrad_gemm( + x: torch.Tensor, + dy: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Perform wgrad GEMM: dw = dy^T * x + + May be fused with bgrad computation. + + May be called outside of this function to enable + some advanced communication/compute overlapping. + + """ + nvtx_range_push(f"{nvtx_label}.wgrad_gemm") + dw, db, *_ = general_gemm(x, dy, **wgrad_gemm_kwargs) + nvtx_range_pop(f"{nvtx_label}.wgrad_gemm") + return dw, db + + # Choose whether to call wgrad GEMM now or delay + if args.wgrad_store is not None and args.wgrad_store.delay_wgrad_compute(): + if ( + wgrad_gemm_kwargs["ub"] is not None + or wgrad_gemm_kwargs["ub_type"] is not None + or wgrad_gemm_kwargs["extra_output"] is not None + or wgrad_gemm_kwargs["bulk_overlap"] + ): + raise NotImplementedError( + "Delayed weight grad computation is not supported " + "with Userbuffers (tensor-parallel communication overlapping)" ) - else: - dgrad, dgrad_work = allreduce(dgrad, ctx.tp_group, async_op=True) - nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_dgrad") + args.wgrad_store.put([ln_out_total, grad_output], wgrad_gemm) else: - dgrad = gemm_out - - # -------------------------------------------------- - # Grad input tensor has been computed... - # -------------------------------------------------- - - # cuBLASMp's AG+GEMM consumes the gathered grad_output inline and - # does not preserve it for wgrad. Userbuffers leaves the gathered - # tensor in its persistent buffer; cuBLASMp does not, so we gather - # here. Route through the same FP8-aware all-gather as the - # non-overlap path in - # ``TransformerEngineBaseModule.grad_output_preprocess`` by passing - # the grad_output quantizer. Columnwise data needed for wgrad is - # produced by ``update_usage(columnwise_usage=True)`` further below. - if ( - ctx.requires_wgrad - and ctx.ub_overlap_ag - and ctx.ub_obj_gradout is not None - and ctx.ub_obj_gradout.with_cublasmp() - ): - if ctx.grad_output_quantizer is not None: - set_quantizer_usage_for_wgrad_all_gather(ctx.grad_output_quantizer) - grad_output, _ = gather_along_first_dim( - grad_output, - ctx.tp_group, - quantizer=ctx.grad_output_quantizer, - ) - - # -------------------------------------------------- - # Compute grad weight - # -------------------------------------------------- - wgrad = None - if ctx.requires_wgrad: - # Prepare grad output tensor - # Note: Synchronize tensor-parallel communication and - # make sure required data is available - if ctx.ub_overlap_ag and isinstance(ctx.grad_output_quantizer, MXFP8Quantizer): - # UB does not support pipelined overlapping grad output - # all-gather with wgrad GEMM. Also, we can't - # convert row-scaled MXFP8 to column-scaled, so we - # can't reuse the grad output that was gathered - # for the dgrad GEMM. We work around by explicitly - # overlapping the AG operation with the dgrad GEMM. - - # Get the communication stream from the dgrad GEMM to use for the AG - dgrad_send_stream, dgrad_recv_stream = ub_obj_dgrad.get_communication_stream() - - # This object is separate from the ub_obj_wgrad object which is passed to the GEMM - ub_obj_overlap_wgrad = get_ub(ctx.ub_name + "_wgrad", ctx.fp8) - - ctx.grad_output_quantizer.set_usage(rowwise=False, columnwise=True) - - # We use the send stream to copy into the userbuffers. - # This is the same stream that we will use to access the data in the AG, - # so we dont need to add any syncs yet. - with torch.cuda.stream(dgrad_send_stream): - grad_output, _ = fill_userbuffers_buffer_for_all_gather( - ub_obj_overlap_wgrad, - grad_outputs[0], - ctx.grad_output_quantizer, - ctx.tp_group, - ) - - # Allgather grad_outputs[0] using the dgrad streams so we can overlap with the fc2_dgrad gemm - tex.bulk_overlap_ag_with_external_gemm( - ub_obj_overlap_wgrad, dgrad_send_stream, dgrad_recv_stream - ) + # Call wgrad GEMM now + wgrad, grad_bias_ = wgrad_gemm(ln_out_total, grad_output) - # Prepare input tensor - # Note: Synchronize tensor-parallel communication and - # make sure required data is available - if ln_out_total_work is not None: - ln_out_total_work.wait() - ln_out_total_work = None - if ctx.fp8 or ctx.debug: - if isinstance(ln_out_total, QuantizedTensorStorage): - ln_out_total.update_usage(columnwise_usage=True) - else: - ctx.input_quantizer.set_usage(rowwise=False, columnwise=True) - ln_out_total = ctx.input_quantizer(ln_out_total) - - if ctx.fp8 or ctx.debug: - if isinstance(grad_output, QuantizedTensorStorage): - grad_output.update_usage(columnwise_usage=True) - else: - ctx.grad_output_quantizer.set_usage(rowwise=False, columnwise=True) - grad_output = ctx.grad_output_quantizer(grad_output) - - # Figure out whether to use split accumulator - use_split_accumulator = _2X_ACC_WGRAD - if ctx.fp8: - recipe = ctx.fp8_recipe - if hasattr(recipe, "fp8_gemm_wgrad"): - use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator - - # Figure out whether to output wgrad GEMM directly into main grad 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 - ) + wgrad = finalize_weight_grads(saved_weight, [wgrad])[0] + + # Update grad bias if needed + if grad_bias is None: + grad_bias = grad_bias_ + del grad_bias_ + + # Deallocate input tensors if permitted + if not args.return_layernorm_output and not args.return_layernorm_output_gathered: + # Input tensors have not been exposed externally + clear_tensor_data(ln_out) + elif args.ln_out_needs_gather and args.return_layernorm_output_gathered: + # Non-gathered input has not been exposed externally + clear_tensor_data(ln_out) + if args.ln_out_needs_gather: + # Gathered input is internal + clear_tensor_data(ln_out_total) + if args.sequence_parallel and ( + args.parallel_mode == "row" or (args.parallel_mode == "column" and args.fp8) + ): + # Gathered (row-SP) or quantized (column-SP FP8) grad_output is internal + clear_tensor_data(grad_output) + + # Update grad input if overlapping reduce-scatter with wgrad GEMM + if args.ub_bulk_wgrad: + if ub_obj_wgrad.is_fp8_ubuf(): + dgrad = reduce_scatter_out else: - accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation - - # Output buffer for overlapping FP8 grad input - # reduce-scatter with wgrad GEMM - reduce_scatter_out = None - if ctx.ub_bulk_wgrad and ub_obj_wgrad.is_fp8_ubuf(): - reduce_scatter_out = torch.empty( - dgrad_shape, dtype=ctx.activation_dtype, device=grad_outputs[0].device - ) + dgrad = ub_obj_wgrad.get_buffer(local_chunk=True).clone() - # Arguments to include in wgrad GEMM closure - wgrad_gemm_kwargs = { - "out_dtype": ( - main_grad.dtype if ctx.fuse_wgrad_accumulation else ctx.activation_dtype - ), - "quantization_params": ctx.grad_weight_quantizer, - "accumulate": ( - accumulate_wgrad_into_param_main_grad - if not origin_weight_overwrites_main_grad - else False - ), - "layout": "NT", - "out": main_grad if ctx.fuse_wgrad_accumulation else None, - "bias": (bias if (grad_bias is None and not ctx.fp8) else None), - "use_split_accumulator": use_split_accumulator, - "grad": True, - "ub": ub_obj_wgrad, - "ub_type": ub_type_wgrad, - "extra_output": reduce_scatter_out, - "bulk_overlap": ctx.ub_bulk_wgrad, - } - - def wgrad_gemm( - x: torch.Tensor, - dy: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Perform wgrad GEMM: dw = dy^T * x - - May be fused with bgrad computation. - - May be called outside of this function to enable - some advanced communication/compute overlapping. - - """ - nvtx_range_push(f"{nvtx_label}.wgrad_gemm") - dw, db, *_ = general_gemm(x, dy, **wgrad_gemm_kwargs) - nvtx_range_pop(f"{nvtx_label}.wgrad_gemm") - return dw, db - - # Choose whether to call wgrad GEMM now or delay - if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute(): - if ( - wgrad_gemm_kwargs["ub"] is not None - or wgrad_gemm_kwargs["ub_type"] is not None - or wgrad_gemm_kwargs["extra_output"] is not None - or wgrad_gemm_kwargs["bulk_overlap"] - ): - raise NotImplementedError( - "Delayed weight grad computation is not supported " - "with Userbuffers (tensor-parallel communication overlapping)" - ) - ctx.wgrad_store.put([ln_out_total, grad_output], wgrad_gemm) - else: + # -------------------------------------------------- + # Grad weight has been computed... + # -------------------------------------------------- - # 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_ - del grad_bias_ - - # Deallocate input tensors if permitted - if not ctx.return_layernorm_output and not ctx.return_layernorm_output_gathered: - # Input tensors have not been exposed externally - clear_tensor_data(ln_out) - elif ctx.ln_out_needs_gather and ctx.return_layernorm_output_gathered: - # Non-gathered input has not been exposed externally - clear_tensor_data(ln_out) - if ctx.ln_out_needs_gather: - # Gathered input is internal - clear_tensor_data(ln_out_total) - if ctx.sequence_parallel and ( - ctx.parallel_mode == "row" or (ctx.parallel_mode == "column" and ctx.fp8) - ): - # Gathered (row-SP) or quantized (column-SP FP8) grad_output is internal - clear_tensor_data(grad_output) - - # Update grad input if overlapping reduce-scatter with wgrad GEMM - if ctx.ub_bulk_wgrad: - if ub_obj_wgrad.is_fp8_ubuf(): - dgrad = reduce_scatter_out - else: - dgrad = ub_obj_wgrad.get_buffer(local_chunk=True).clone() - - # -------------------------------------------------- - # Grad weight has been computed... - # -------------------------------------------------- - - # Don't return grad bias if not needed - if not ctx.use_bias: - grad_bias = None - - # Synchronize tensor parallel communication - if ln_out_total_work is not None: - ln_out_total_work.wait() - ln_out_total_work = None - if dgrad_work is not None: - dgrad_work.wait() - dgrad_work = None + # Don't return grad bias if not needed + if not args.use_bias: + grad_bias = None + + # Synchronize tensor parallel communication + if ln_out_total_work is not None: + ln_out_total_work.wait() + ln_out_total_work = None + if dgrad_work is not None: + dgrad_work.wait() + dgrad_work = None - # Residual gradient - dgrad = dgrad.view(inputmat.shape) - if ctx.return_layernorm_output and not ctx.return_layernorm_output_gathered: - dgrad = dgrad + grad_outputs[1].view_as(dgrad) + # Residual gradient + dgrad = dgrad.view(inputmat.shape) + if args.return_layernorm_output and not args.return_layernorm_output_gathered: + dgrad = dgrad + args.grad_ln_out.view_as(dgrad) - # Norm gradient - dgamma = None + # Norm gradient + dgamma = None + dbeta = None + nvtx_range_push(f"{nvtx_label}.norm") + if args.normalization == "LayerNorm": + dgrad, dgamma, dbeta = tex.layernorm_bwd( + dgrad, + inputmat, + mu, + rsigma, + ln_weight, + args.bwd_ln_sm_margin, + args.zero_centered_gamma, + ) + dgrad = dgrad.reshape(inputmat.size()) + elif args.normalization == "RMSNorm": + dgrad, dgamma = tex.rmsnorm_bwd( + dgrad, + inputmat, + rsigma, + ln_weight, + args.bwd_ln_sm_margin, + args.zero_centered_gamma, + ) + dgrad = dgrad.reshape(inputmat.size()) dbeta = None - nvtx_range_push(f"{nvtx_label}.norm") - if ctx.normalization == "LayerNorm": - dgrad, dgamma, dbeta = tex.layernorm_bwd( - dgrad, - inputmat, - mu, - rsigma, - ln_weight, - ctx.bwd_ln_sm_margin, - ctx.zero_centered_gamma, + nvtx_range_pop(f"{nvtx_label}.norm") + clear_tensor_data(mu) + clear_tensor_data(rsigma) + + if args.requires_wgrad: + # Handle custom DDP from mcore. + if args.fuse_wgrad_accumulation and hasattr(origin_weight, "grad_added_to_main_grad"): + origin_weight.grad_added_to_main_grad = True + if getattr(origin_weight, "zero_out_wgrad", False): + wgrad = get_dummy_wgrad( + list(main_grad.shape), + origin_weight.dtype, + zero=True, ) - dgrad = dgrad.reshape(inputmat.size()) - elif ctx.normalization == "RMSNorm": - dgrad, dgamma = tex.rmsnorm_bwd( - dgrad, - inputmat, - rsigma, - ln_weight, - ctx.bwd_ln_sm_margin, - ctx.zero_centered_gamma, + else: + wgrad = get_dummy_wgrad( + list(main_grad.shape), + origin_weight.dtype, ) - dgrad = dgrad.reshape(inputmat.size()) - dbeta = None - nvtx_range_pop(f"{nvtx_label}.norm") - clear_tensor_data(mu) - clear_tensor_data(rsigma) - - if ctx.requires_wgrad: - # Handle custom DDP from mcore. - if ctx.fuse_wgrad_accumulation and hasattr(origin_weight, "grad_added_to_main_grad"): - origin_weight.grad_added_to_main_grad = True - if getattr(origin_weight, "zero_out_wgrad", False): - wgrad = get_dummy_wgrad( - list(main_grad.shape), - origin_weight.dtype, - zero=True, - ) - else: - wgrad = get_dummy_wgrad( - list(main_grad.shape), - origin_weight.dtype, - ) - elif ctx.fuse_wgrad_accumulation: - wgrad = None - else: + elif args.fuse_wgrad_accumulation: wgrad = None + else: + wgrad = None + + return ( + dgrad.view(args.inp_shape) if args.requires_dgrad else None, + dgamma, + dbeta, + wgrad, + grad_bias, + ) + + +class _LayerNormLinear(torch.autograd.Function): + """LayerNormLinear semi-top level module + Calls custom cuda extensions. + """ - if ctx.reduce_and_update_bwd_fp8_tensors and not is_graph_capturing(): + @staticmethod + def forward( + ctx, + inp: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: Optional[torch.Tensor], + weight: torch.Tensor, + bias: Optional[torch.Tensor], + fwd_args: LayerNormLinearFwdArgs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: + """Forward pass: compute the output and set up the autograd context. + + The tensors are positional so autograd tracks them; they are + re-attached to ``fwd_args`` so every downstream helper takes a single + argument. ``weight_workspace`` is a non-differentiable cached tensor + passed in via ``fwd_args`` and the freshly produced workspace is + returned as a separate output so the module can refresh its cache. + """ + fwd_args.inp = inp + fwd_args.ln_weight = ln_weight + fwd_args.ln_bias = ln_bias + fwd_args.weight = weight + fwd_args.bias = bias + ( + out, + ln_out_return, + new_weight_workspace, + tensors_to_save_from_forward, + ctx_attrs, + ) = _layernorm_linear_forward_impl(fwd_args) + if ctx is not None: + bwd_args = LayerNormLinearBwdArgs() + tensors_to_save_from_setup = _layernorm_linear_setup_ctx( + bwd_args, + fwd_args, + (out, ln_out_return, new_weight_workspace), + ctx_attrs, + tensors_to_save_from_forward, + ) + tensors_to_save, tensor_objects = prepare_for_saving(*tensors_to_save_from_setup) + ctx.save_for_backward(*tensors_to_save) + ctx.tensor_objects = tensor_objects + ctx.backward_objects = bwd_args + if fwd_args.fp8 and fwd_args.any_requires_grad(): + bwd_args.reduce_and_update_bwd_fp8_tensors = check_fp8_reduce_and_update() + if fwd_args.backward_override is not None: + bwd_args.reduce_and_update_bwd_fp8_tensors = False + + return out, ln_out_return, new_weight_workspace + + @staticmethod + def backward( + ctx, + grad_output: torch.Tensor, + grad_ln_out: Optional[torch.Tensor], + _grad_weight_workspace, + ) -> Tuple[Union[torch.Tensor, None], ...]: + """Backward pass: compute gradients and reduce FP8 scaling factors.""" + bwd_args: LayerNormLinearBwdArgs = ctx.backward_objects + bwd_args.grad_output = grad_output + bwd_args.grad_ln_out = grad_ln_out + bwd_args.setup_saved_tensors(ctx) + nvtx_label = "transformer_engine._LayerNormLinear.backward" + if bwd_args.ub_name is not None: + nvtx_label = f"{nvtx_label}.{bwd_args.ub_name}" + dgrad, dgamma, dbeta, wgrad, grad_bias = _layernorm_linear_backward_impl(bwd_args) + reduce_and_update_bwd_fp8_tensors = bwd_args.reduce_and_update_bwd_fp8_tensors + # Drop all references held by bwd_args (saved tensors, quantizers, weakrefs, + # main_grad closure) so they don't outlive backward via ctx under retain_graph. + ctx.backward_objects = None + del bwd_args + if reduce_and_update_bwd_fp8_tensors and not is_graph_capturing(): nvtx_range_push(f"{nvtx_label}.reduce_and_update_fp8_tensors") FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) nvtx_range_pop(f"{nvtx_label}.reduce_and_update_fp8_tensors") - - # Scatter fp8 weight buffers - # if ctx.fp8 and not isinstance(weight, QuantizedTensorStorage): - # _fsdp_scatter_tensors(ctx.fsdp_group, weight_fp8) - return ( - dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, + dgrad, dgamma, dbeta, wgrad, - None, # weight_workspace grad_bias, - None, + None, # fwd_args ) @@ -1754,69 +2105,136 @@ def forward( weight_quantizer, weight_tensor ) - if is_grad_enabled: - fwd_fn = _LayerNormLinear.apply - autograd_ctx = [] - else: - fwd_fn = _LayerNormLinear.forward - autograd_ctx = [None] cache_name = None if (is_first_microbatch is None or self.is_fsdp2) else "weight" weight_workspace = ( self._fp8_workspaces.get(cache_name) if cache_name is not None else None ) - non_tensor_args = ( - self.eps, - is_first_microbatch, - self.fp8, - self.fp8_calibration, - self.wgrad_store, - self.fuse_wgrad_accumulation, - input_quantizer, - weight_quantizer, - output_quantizer, - grad_input_quantizer, - grad_weight_quantizer, - grad_output_quantizer, - is_cpu_offload_enabled(), - self.tp_group, - self.tp_size, - self.sequence_parallel, - self.tp_size > 1, - self.activation_dtype, - self.parallel_mode, - self.return_layernorm_output, - self.return_layernorm_output_gathered, - is_grad_enabled, - self.fwd_ln_sm_margin if is_grad_enabled else self.inf_ln_sm_margin, - self.bwd_ln_sm_margin, - self.zero_centered_gamma, - self.normalization, - self.ub_overlap_ag_fprop, - self.ub_overlap_rs_fprop, - self.ub_overlap_ag_dgrad, - self.ub_overlap_rs_dgrad, - self.ub_bulk_wgrad, - self.ub_bulk_dgrad, - self.ub_name, - self.fsdp_group, - cache_name is not None, - skip_fp8_weight_update, - self.symmetric_ar_type, - debug, - self.is_fsdp2, + dgrad_use_split_accumulator = _2X_ACC_DGRAD + wgrad_use_split_accumulator = _2X_ACC_WGRAD + if self.fp8: + _recipe = FP8GlobalStateManager.get_fp8_recipe() + backward_override = _recipe.backward_override + if hasattr(_recipe, "fp8_gemm_dgrad"): + dgrad_use_split_accumulator = _recipe.fp8_gemm_dgrad.use_split_accumulator + if hasattr(_recipe, "fp8_gemm_wgrad"): + wgrad_use_split_accumulator = _recipe.fp8_gemm_wgrad.use_split_accumulator + else: + backward_override = None + + if debug: # turn off userbuffers in debug mode + ub_overlap_ag_fprop = False + ub_overlap_rs_fprop = False + ub_overlap_ag_dgrad = False + ub_overlap_rs_dgrad = False + ub_bulk_wgrad = False + ub_bulk_dgrad = False + else: + ub_overlap_ag_fprop = self.ub_overlap_ag_fprop + ub_overlap_rs_fprop = self.ub_overlap_rs_fprop + ub_overlap_ag_dgrad = self.ub_overlap_ag_dgrad + ub_overlap_rs_dgrad = self.ub_overlap_rs_dgrad + ub_bulk_wgrad = self.ub_bulk_wgrad + ub_bulk_dgrad = self.ub_bulk_dgrad + + linear_bias_tensor = ( + bias_tensor if (self.apply_bias and not self.gemm_bias_unfused_add) else None ) - out, ln_out, new_weight_workspace = fwd_fn( - *autograd_ctx, - inp, - self.layer_norm_weight, - self.layer_norm_bias, - weight_tensor, - weight_workspace, - bias_tensor if self.apply_bias and not self.gemm_bias_unfused_add else None, - non_tensor_args, + fwd_args = LayerNormLinearFwdArgs( + # tensors + inp=inp, + ln_weight=self.layer_norm_weight, + ln_bias=self.layer_norm_bias, + weight=weight_tensor, + bias=linear_bias_tensor, + weight_workspace=weight_workspace, + # requires_grad flags + input_requires_grad=inp.requires_grad, + ln_weight_requires_grad=self.layer_norm_weight.requires_grad, + ln_bias_requires_grad=( + self.layer_norm_bias.requires_grad + if self.layer_norm_bias is not None + else False + ), + weight_requires_grad=weight_tensor.requires_grad, + bias_requires_grad=( + linear_bias_tensor.requires_grad if linear_bias_tensor is not None else False + ), + # quantizers + input_quantizer=input_quantizer, + weight_quantizer=weight_quantizer, + output_quantizer=output_quantizer, + grad_input_quantizer=grad_input_quantizer, + grad_weight_quantizer=grad_weight_quantizer, + grad_output_quantizer=grad_output_quantizer, + # normalization + eps=self.eps, + normalization=self.normalization, + zero_centered_gamma=self.zero_centered_gamma, + fwd_ln_sm_margin=( + self.fwd_ln_sm_margin if is_grad_enabled else self.inf_ln_sm_margin + ), + bwd_ln_sm_margin=self.bwd_ln_sm_margin, + return_layernorm_output=self.return_layernorm_output, + return_layernorm_output_gathered=self.return_layernorm_output_gathered, + # numerical / dtype config + activation_dtype=self.activation_dtype, + fp8=self.fp8, + fp8_calibration=self.fp8_calibration, + backward_override=backward_override, + dgrad_use_split_accumulator=dgrad_use_split_accumulator, + wgrad_use_split_accumulator=wgrad_use_split_accumulator, + debug=debug, + # weight-workspace caching + is_first_microbatch=is_first_microbatch, + cache_weight=cache_name is not None, + skip_fp8_weight_update=skip_fp8_weight_update, + # tensor / sequence parallelism + parallel_mode=self.parallel_mode, + tp_group=self.tp_group, + tp_size=self.tp_size, + tensor_parallel=self.tp_size > 1, + sequence_parallel=self.sequence_parallel, + symmetric_ar_type=self.symmetric_ar_type, + # userbuffers + ub_name=self.ub_name, + ub_overlap_ag_fprop=ub_overlap_ag_fprop, + ub_overlap_rs_fprop=ub_overlap_rs_fprop, + ub_overlap_ag_dgrad=ub_overlap_ag_dgrad, + ub_overlap_rs_dgrad=ub_overlap_rs_dgrad, + ub_bulk_dgrad=ub_bulk_dgrad, + ub_bulk_wgrad=ub_bulk_wgrad, + # FSDP + fsdp_group=self.fsdp_group, + is_fsdp2=self.is_fsdp2, + # weight-grad scheduling + fuse_wgrad_accumulation=self.fuse_wgrad_accumulation, + wgrad_store=self.wgrad_store, + # misc + cpu_offloading=is_cpu_offload_enabled(), + is_grad_enabled=is_grad_enabled, ) + if is_grad_enabled: + out, ln_out, new_weight_workspace = _LayerNormLinear.apply( + inp, + self.layer_norm_weight, + self.layer_norm_bias, + weight_tensor, + linear_bias_tensor, + fwd_args, + ) + else: + out, ln_out, new_weight_workspace = _LayerNormLinear.forward( + None, + inp, + self.layer_norm_weight, + self.layer_norm_bias, + weight_tensor, + linear_bias_tensor, + fwd_args, + ) + if new_weight_workspace is not None and cache_name is not None: if isinstance(new_weight_workspace, torch.Tensor): new_weight_workspace = new_weight_workspace.detach() diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 3ee0cda50c5..8d2dc192159 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -5,8 +5,9 @@ """LayerNormMLP API""" import os import warnings +from dataclasses import dataclass, replace as dataclass_replace import weakref -from typing import Callable, Optional, Tuple, Union, List +from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union, List from functools import reduce from operator import mul as multiply_op @@ -46,7 +47,6 @@ cast_if_needed, assert_dim_for_fp8_exec, clear_tensor_data, - requires_grad, needs_quantized_gemm, get_nvtx_range_context, ) @@ -58,7 +58,6 @@ reduce_scatter_along_first_dim, gather_along_first_dim, use_reentrant_activation_recompute, - in_fp8_activation_recompute_phase, _fsdp_scatter_tensors, _get_cuda_rng_state, _set_cuda_rng_state, @@ -74,6 +73,7 @@ from ..tensor.identity_tensor import IdentityQuantizer from ._common import ( apply_normalization, + check_fp8_reduce_and_update, set_quantizer_amax_reduction_group, set_quantizer_usage_for_wgrad_all_gather, WeightGradStore, @@ -90,6 +90,7 @@ prepare_for_saving, restore_from_func_ctx, ) +from ..dynamo import TensorOrQuantized from ..cpp_extensions import ( general_gemm, ) @@ -171,1663 +172,2007 @@ def _act_func(activation: str, recipe: Optional[Recipe] = None): return funcs[activation] -class _LayerNormMLP(torch.autograd.Function): - """LayerNormMLP semi-top level module - Calls custom cuda extensions. - """ - - @staticmethod - def _forward( - ctx, - inp: torch.Tensor, - ln_weight: torch.Tensor, - ln_bias: torch.Tensor, - fc1_weight: torch.Tensor, - fc1_weight_workspace: Optional[torch.Tensor], - fc1_bias: torch.Tensor, - fc2_weight: torch.Tensor, - fc2_weight_workspace: Optional[torch.Tensor], - fc2_bias: torch.Tensor, - non_tensor_args: Tuple, - ) -> Union[Tuple[torch.Tensor, ...], torch.Tensor]: - # pylint: disable=missing-function-docstring - - # Reduce number of arguments to autograd function in order - # to reduce CPU overhead due to pytorch arg checking. - ( - eps, - is_first_microbatch, - fp8, - fp8_calibration, - wgrad_store, - fuse_wgrad_accumulation, - fc1_input_quantizer, - fc1_weight_quantizer, - fc1_output_quantizer, - fc1_grad_input_quantizer, - fc1_grad_weight_quantizer, - fc1_grad_output_quantizer, - fc2_input_quantizer, - fc2_weight_quantizer, - fc2_output_quantizer, - fc2_grad_input_quantizer, - fc2_grad_weight_quantizer, - fc2_grad_output_quantizer, - cpu_offloading, - tp_group, - tp_size, - sequence_parallel, - tensor_parallel, - activation_dtype, - return_layernorm_output, - return_layernorm_output_gathered, - bias_gelu_fusion, - set_parallel_mode, - is_grad_enabled, - fwd_ln_sm_margin, - bwd_ln_sm_margin, - zero_centered_gamma, - activation, - activation_params, - normalization, - ub_overlap_ag, - ub_overlap_rs, - ub_overlap_rs_dgrad, - ub_bulk_wgrad, - ub_bulk_dgrad, - gemm_gelu_fusion, - fsdp_group, - fp8_meta, - cache_weight, - skip_fp8_weight_update, - symmetric_ar_type, - checkpoint, - debug, - is_fsdp2, - recompute_for_bwd, - ) = non_tensor_args - if fp8: - backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override - else: - backward_override = None - assert backward_override is None, ( - "NVTE_BACKWARD_OVERRIDE=high_precision/dequantized is not implemented in LayerNormMLP." - " Replace LayerNormMLP with LayerNormLinear + Linear to enable" - " high_precision/dequantized backward." +@dataclass(slots=True) +class LayerNormMLPFwdArgs: + """Single-argument bag for the forward path of :class:`_LayerNormMLP`.""" + + # --- Differentiable tensors (also passed positionally to autograd) --- + inp: torch.Tensor + ln_weight: torch.Tensor + ln_bias: Optional[torch.Tensor] + fc1_weight: TensorOrQuantized + fc1_bias: Optional[torch.Tensor] + fc2_weight: TensorOrQuantized + fc2_bias: Optional[torch.Tensor] + + # --- Non-differentiable cached tensors --- + fc1_weight_workspace: Optional[TensorOrQuantized] + fc2_weight_workspace: Optional[TensorOrQuantized] + + # --- requires_grad flags (cached so backward does not re-query) --- + input_requires_grad: bool + ln_weight_requires_grad: bool + ln_bias_requires_grad: bool + fc1_weight_requires_grad: bool + fc1_bias_requires_grad: bool + fc2_weight_requires_grad: bool + fc2_bias_requires_grad: bool + + # --- Quantizers --- + fc1_input_quantizer: Optional[Quantizer] + fc1_weight_quantizer: Optional[Quantizer] + fc1_output_quantizer: Optional[Quantizer] + fc1_grad_input_quantizer: Optional[Quantizer] + fc1_grad_weight_quantizer: Optional[Quantizer] + fc1_grad_output_quantizer: Optional[Quantizer] + fc2_input_quantizer: Optional[Quantizer] + fc2_weight_quantizer: Optional[Quantizer] + fc2_output_quantizer: Optional[Quantizer] + fc2_grad_input_quantizer: Optional[Quantizer] + fc2_grad_weight_quantizer: Optional[Quantizer] + fc2_grad_output_quantizer: Optional[Quantizer] + + # --- Normalization --- + eps: float + normalization: str + zero_centered_gamma: bool + fwd_ln_sm_margin: int + bwd_ln_sm_margin: int + return_layernorm_output: bool + return_layernorm_output_gathered: bool + + # --- Activation --- + activation: str + activation_params: Optional[Dict[str, Any]] + bias_gelu_fusion: bool + gemm_gelu_fusion: bool + + # --- Numerical / dtype config --- + activation_dtype: torch.dtype + fp8: bool + fp8_calibration: bool + backward_override: Optional[str] + dgrad_use_split_accumulator: bool + wgrad_use_split_accumulator: bool + debug: bool + + # --- Weight-workspace caching --- + is_first_microbatch: Optional[bool] + cache_weight: bool + skip_fp8_weight_update: Optional[torch.Tensor] + + # --- Tensor / sequence parallelism --- + set_parallel_mode: bool + tp_group: Optional[dist_group_type] + tp_size: int + tensor_parallel: bool + sequence_parallel: bool + symmetric_ar_type: Optional[str] + + # --- Userbuffers (comm + GEMM overlap) --- + ub_overlap_ag: bool + ub_overlap_rs: bool + ub_overlap_rs_dgrad: bool + ub_bulk_dgrad: bool + ub_bulk_wgrad: bool + + # --- FSDP --- + fsdp_group: Optional[Any] + is_fsdp2: bool + + # --- Weight-grad scheduling --- + fuse_wgrad_accumulation: bool + wgrad_store: Optional[Any] + + # --- Activation checkpointing (recompute in backward) --- + checkpoint: bool + fp8_meta: Optional[Any] + recompute_for_bwd: bool + + # --- Misc --- + cpu_offloading: bool + is_grad_enabled: bool + + def any_requires_grad(self) -> bool: + """Whether any differentiable input requires a gradient.""" + return any( + ( + self.input_requires_grad, + self.ln_weight_requires_grad, + self.ln_bias_requires_grad, + self.fc1_weight_requires_grad, + self.fc1_bias_requires_grad, + self.fc2_weight_requires_grad, + self.fc2_bias_requires_grad, + ) ) - # if grad is enabled and this is not the bwd stage, we must save this so bwd knows which path to take - if is_grad_enabled and not recompute_for_bwd: - ctx.checkpoint = checkpoint - if checkpoint: - # save the state of autocast and quantizers for recomputation - ctx.autocast_state = FP8GlobalStateManager.get_autocast_state() - if ( - fp8 - and FP8GlobalStateManager.get_fp8_recipe().__class__.__name__ - == "DelayedScaling" - ): # only applicable for delayed scaling - FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute( - fp8_meta - ) # to restore quantizers during recomputation - # save the rng states - ctx.cpu_rng_state = torch.get_rng_state() - ctx.cuda_rng_state = _get_cuda_rng_state() - - # whether to save activations regularly, or save inputs for recomputation in bwd - save_for_checkpoint = checkpoint and is_grad_enabled and not recompute_for_bwd - - # whether we are in the forward stage, or recomputing in the bwd stage (false if not checkpointing) - is_recomputation = checkpoint and is_grad_enabled and recompute_for_bwd - - # save the initial state for recomputation by bwd - if save_for_checkpoint: - - # save tensors - tensors_to_save, tensor_objects = prepare_for_saving( - inp, - ln_weight, - ln_bias, - fc1_weight, - fc1_bias, - fc2_weight, - fc2_bias, - ) - ctx.save_for_backward(*tensors_to_save) - ctx.tensor_objects = tensor_objects - ctx.other_args = { - "eps": eps, - "is_first_microbatch": is_first_microbatch, - "fp8": fp8, - "fp8_calibration": fp8_calibration, - "wgrad_store": wgrad_store, - "fuse_wgrad_accumulation": fuse_wgrad_accumulation, - "fc1_input_quantizer": fc1_input_quantizer, - "fc1_weight_quantizer": fc1_weight_quantizer, - "fc1_output_quantizer": fc1_output_quantizer, - "fc1_grad_input_quantizer": fc1_grad_input_quantizer, - "fc1_grad_weight_quantizer": fc1_grad_weight_quantizer, - "fc1_grad_output_quantizer": fc1_grad_output_quantizer, - "fc2_input_quantizer": fc2_input_quantizer, - "fc2_weight_quantizer": fc2_weight_quantizer, - "fc2_output_quantizer": fc2_output_quantizer, - "fc2_grad_input_quantizer": fc2_grad_input_quantizer, - "fc2_grad_weight_quantizer": fc2_grad_weight_quantizer, - "fc2_grad_output_quantizer": fc2_grad_output_quantizer, - "cpu_offloading": cpu_offloading, - "tp_group": tp_group, - "tp_size": tp_size, - "sequence_parallel": sequence_parallel, - "tensor_parallel": tensor_parallel, - "activation_dtype": activation_dtype, - "return_layernorm_output": return_layernorm_output, - "return_layernorm_output_gathered": return_layernorm_output_gathered, - "bias_gelu_fusion": bias_gelu_fusion, - "set_parallel_mode": set_parallel_mode, - "is_grad_enabled": is_grad_enabled, - "fwd_ln_sm_margin": fwd_ln_sm_margin, - "bwd_ln_sm_margin": bwd_ln_sm_margin, - "zero_centered_gamma": zero_centered_gamma, - "activation": activation, - "activation_params": activation_params, - "normalization": normalization, - "ub_overlap_ag": ub_overlap_ag, - "ub_overlap_rs": ub_overlap_rs, - "ub_overlap_rs_dgrad": ub_overlap_rs_dgrad, - "ub_bulk_wgrad": ub_bulk_wgrad, - "ub_bulk_dgrad": ub_bulk_dgrad, - "gemm_gelu_fusion": gemm_gelu_fusion, - "fsdp_group": fsdp_group, - "fp8_meta": fp8_meta, - "cache_weight": False, - "skip_fp8_weight_update": skip_fp8_weight_update, - "symmetric_ar_type": symmetric_ar_type, - "checkpoint": checkpoint, - "debug": debug, - "is_fsdp2": is_fsdp2, - "recompute_for_bwd": True, # set this to true for recomputation phase - } - # Make sure input dimensions are compatible - in_features, inp_shape = ln_weight.numel(), inp.shape - assert inp_shape[-1] == in_features, "GEMM not possible" - inputmat = inp.view((-1, in_features)) - if fp8: - assert_dim_for_fp8_exec(inputmat, fc1_weight, fc2_weight) - - activation_func = _act_func( - activation, FP8GlobalStateManager.get_fp8_recipe() if fp8 else None - )[0] - - # Cast for native AMP - inputmat = cast_if_needed(inputmat, activation_dtype) - ln_weight = cast_if_needed(ln_weight, activation_dtype) - if ln_bias is not None: - ln_bias = cast_if_needed(ln_bias, activation_dtype) - if is_cpu_offload_enabled(): - start_offload(inputmat) - - tp_world_size = get_distributed_world_size(tp_group) - - # bwd needs fc1 input when grad is enabled, fc1 needs grad, and either - # 1) no checkpointing - # or 2) doing the recomputation with checkpointing - backwards_needs_fc1_input = fc1_weight.requires_grad and ( - (is_grad_enabled and not checkpoint) or is_recomputation - ) +@dataclass(slots=True) +class LayerNormMLPBwdArgs: + """Single-argument bag for the backward path of :class:`_LayerNormMLP`.""" + + # --- Incoming gradients (populated at backward entry) --- + grad_output: Optional[torch.Tensor] = None + grad_ln_out: Optional[torch.Tensor] = None + + # --- Saved / restored tensors (populated at backward entry) --- + inputmat: Optional[torch.Tensor] = None + ln_weight: Optional[torch.Tensor] = None + ln_out: Optional[TensorOrQuantized] = None + fc1_weight_fp8: Optional[TensorOrQuantized] = None + fc1_weight: Optional[TensorOrQuantized] = None + fc1_bias: Optional[torch.Tensor] = None + fc1_out: Optional[TensorOrQuantized] = None + fc1_out_without_bias: Optional[torch.Tensor] = None + act_out: Optional[TensorOrQuantized] = None + fc2_weight_fp8: Optional[TensorOrQuantized] = None + fc2_weight: Optional[TensorOrQuantized] = None + fc2_bias: Optional[torch.Tensor] = None + mu: Optional[torch.Tensor] = None + rsigma: Optional[torch.Tensor] = None + + # --- Activation checkpointing (forward inputs saved, recomputed in backward) --- + checkpoint: bool = False + checkpoint_fwd_args: Optional[Any] = None + autocast_state: Optional[Any] = None + cpu_rng_state: Optional[Any] = None + cuda_rng_state: Optional[Any] = None + + # --- Quantizers --- + fc1_input_quantizer: Optional[Quantizer] = None + fc1_weight_quantizer: Optional[Quantizer] = None + fc1_grad_input_quantizer: Optional[Quantizer] = None + fc1_grad_weight_quantizer: Optional[Quantizer] = None + fc1_grad_output_quantizer: Optional[Quantizer] = None + fc2_input_quantizer: Optional[Quantizer] = None + fc2_weight_quantizer: Optional[Quantizer] = None + fc2_grad_input_quantizer: Optional[Quantizer] = None + fc2_grad_weight_quantizer: Optional[Quantizer] = None + fc2_grad_output_quantizer: Optional[Quantizer] = None + + # --- Differentiability summary --- + use_bias: bool = False + requires_dgrad: bool = False + fc1_weight_requires_grad: bool = False + fc1_bias_requires_grad: bool = False + fc2_weight_requires_grad: bool = False + inp_shape: Optional[torch.Size] = None + + # --- Normalization --- + normalization: str = "LayerNorm" + zero_centered_gamma: bool = False + bwd_ln_sm_margin: int = 0 + return_layernorm_output: bool = False + return_layernorm_output_gathered: bool = False + + # --- Activation --- + activation: str = "gelu" + activation_params: Optional[Dict[str, Any]] = None + bias_gelu_fusion: bool = False + + # --- Numerical / dtype config --- + activation_dtype: Optional[torch.dtype] = None + fp8: bool = False + fp8_recipe: Optional[Any] = None + dgrad_use_split_accumulator: bool = _2X_ACC_DGRAD + wgrad_use_split_accumulator: bool = _2X_ACC_WGRAD + backward_override: Optional[str] = None + debug: bool = False + + # --- Tensor / sequence parallelism --- + set_parallel_mode: bool = False + tp_group: Optional[dist_group_type] = None + tp_size: int = 1 + tensor_parallel: bool = False + sequence_parallel: bool = False + + # --- Userbuffers (comm + GEMM overlap) --- + ub_overlap_ag: bool = False + ub_overlap_rs_dgrad: bool = False + ub_bulk_dgrad: bool = False + ub_bulk_wgrad: bool = False + + # --- FSDP --- + fsdp_group: Optional[Any] = None + fsdp_shapes: Any = None + is_fsdp2: bool = False + + # --- Weight-grad scheduling / accumulation --- + is_first_microbatch: Optional[bool] = None + fuse_wgrad_accumulation: bool = False + wgrad_store: Optional[Any] = None + fc1_weight_ref: Optional[Any] = None + fc2_weight_ref: Optional[Any] = None + fc1_weight_overwrites_main_grad: bool = False + fc2_weight_overwrites_main_grad: bool = False + fc1_main_grad_func: Optional[Callable[[], torch.Tensor]] = None + fc2_main_grad_func: Optional[Callable[[], torch.Tensor]] = None + + # --- FP8 reduce-and-update bookkeeping --- + reduce_and_update_bwd_fp8_tensors: bool = False + + # --- Misc --- + cpu_offloading: bool = False + + # --- Per-backward scratch state (populated inside the backward impl) --- + ub_obj_gradout: Optional[Any] = None + + def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None: + """Pull saved tensors from ``ctx`` into the fields backward consumes.""" + self.set_saved_tensors(restore_from_func_ctx(ctx)) + + def set_saved_tensors(self, tensors: Sequence[Any]) -> None: + """Bind the (restored or recomputed) saved tensors to their fields.""" + ( + self.inputmat, + self.ln_weight, + self.ln_out, + self.fc1_weight_fp8, + self.fc1_weight, + self.fc1_bias, + self.fc1_out, + self.fc1_out_without_bias, + self.act_out, + self.fc2_weight_fp8, + self.fc2_weight, + self.fc2_bias, + self.mu, + self.rsigma, + ) = tensors + + +_CHECKPOINT_SAVED_ALIASES = ( + "inp", + "ln_weight", + "ln_bias", + "fc1_weight", + "fc1_bias", + "fc2_weight", + "fc2_bias", +) - device = inp.device - # Configure Userbuffers communication (comm+GEMM overlap) - if debug: # turn off userbuffers in debug mode - ub_overlap_ag = False - ub_overlap_rs = False - ub_overlap_rs_dgrad = False - ub_bulk_wgrad = False - ub_bulk_dgrad = False - ub_overlap_ag = ub_overlap_ag and is_grad_enabled and not return_layernorm_output_gathered - ub_overlap_rs = ub_overlap_rs and is_grad_enabled +def _layernorm_mlp_forward_impl( + args: LayerNormMLPFwdArgs, +) -> Tuple[ + Optional[torch.Tensor], + Optional[torch.Tensor], + Optional[torch.Tensor], + Optional[torch.Tensor], + Optional[Tuple], + Optional[Dict], +]: + """Forward implementation for the layernorm-MLP layer. + + Returns ``(fc2_out, ln_out_return, new_fc1_weight_workspace, + new_fc2_weight_workspace, tensors_to_save_from_forward, ctx_attrs)``. The + new workspaces are the freshly produced FP8 weight workspaces (returned so + the caller can refresh its cache). The last two are ``None`` when + gradients are disabled. + + With ``args.checkpoint`` the forward saves only its inputs and the + backward recomputes the activations by calling this again with + ``args.recompute_for_bwd``; that call returns ``None`` user outputs and + the recomputed saved tensors. + """ + inp = args.inp + ln_weight = args.ln_weight + ln_bias = args.ln_bias + fc1_weight = args.fc1_weight + fc1_bias = args.fc1_bias + fc2_weight = args.fc2_weight + fc2_bias = args.fc2_bias + fc1_input_quantizer = args.fc1_input_quantizer + fc1_weight_quantizer = args.fc1_weight_quantizer + fc1_output_quantizer = args.fc1_output_quantizer + fc2_input_quantizer = args.fc2_input_quantizer + fc2_weight_quantizer = args.fc2_weight_quantizer + fc2_output_quantizer = args.fc2_output_quantizer + is_first_microbatch = args.is_first_microbatch + fp8 = args.fp8 + fp8_calibration = args.fp8_calibration + debug = args.debug + cpu_offloading = args.cpu_offloading + tp_group = args.tp_group + tp_size = args.tp_size + sequence_parallel = args.sequence_parallel + tensor_parallel = args.tensor_parallel + set_parallel_mode = args.set_parallel_mode + activation_dtype = args.activation_dtype + is_grad_enabled = args.is_grad_enabled + return_layernorm_output = args.return_layernorm_output + return_layernorm_output_gathered = args.return_layernorm_output_gathered + activation = args.activation + bias_gelu_fusion = args.bias_gelu_fusion + gemm_gelu_fusion = args.gemm_gelu_fusion + ub_overlap_ag = args.ub_overlap_ag + ub_overlap_rs = args.ub_overlap_rs + fsdp_group = args.fsdp_group + is_fsdp2 = args.is_fsdp2 + checkpoint = args.checkpoint + recompute_for_bwd = args.recompute_for_bwd + fc1_weight_requires_grad = args.fc1_weight_requires_grad + fc2_weight_requires_grad = args.fc2_weight_requires_grad + + assert args.backward_override is None, ( + "NVTE_BACKWARD_OVERRIDE=high_precision/dequantized is not implemented in LayerNormMLP." + " Replace LayerNormMLP with LayerNormLinear + Linear to enable" + " high_precision/dequantized backward." + ) + + ctx_attrs = None + # if grad is enabled and this is not the bwd stage, we must save this so bwd knows which path to take + if is_grad_enabled and not recompute_for_bwd: + ctx_attrs = {"checkpoint": checkpoint} + if checkpoint: + # save the state of autocast and quantizers for recomputation + ctx_attrs["autocast_state"] = FP8GlobalStateManager.get_autocast_state() + if ( + fp8 + and FP8GlobalStateManager.get_fp8_recipe().__class__.__name__ == "DelayedScaling" + ): # only applicable for delayed scaling + FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute( + args.fp8_meta + ) # to restore quantizers during recomputation + # save the rng states + ctx_attrs["cpu_rng_state"] = torch.get_rng_state() + ctx_attrs["cuda_rng_state"] = _get_cuda_rng_state() + + # whether to save activations regularly, or save inputs for recomputation in bwd + save_for_checkpoint = checkpoint and is_grad_enabled and not recompute_for_bwd + + # whether we are in the forward stage, or recomputing in the bwd stage (false if not checkpointing) + is_recomputation = checkpoint and is_grad_enabled and recompute_for_bwd + + tensors_to_save_from_forward = None + # save the initial state for recomputation by bwd + if save_for_checkpoint: + tensors_to_save_from_forward = (None,) * len(_CHECKPOINT_SAVED_ALIASES) + ctx_attrs["saved_tensor_aliases"] = _CHECKPOINT_SAVED_ALIASES + + # Make sure input dimensions are compatible + in_features, inp_shape = ln_weight.numel(), inp.shape + assert inp_shape[-1] == in_features, "GEMM not possible" + inp = inp.view((-1, in_features)) + inputmat = inp + if fp8: + assert_dim_for_fp8_exec(inputmat, fc1_weight, fc2_weight) + + activation_func = _act_func( + activation, FP8GlobalStateManager.get_fp8_recipe() if fp8 else None + )[0] + + # Cast for native AMP + inputmat = cast_if_needed(inputmat, activation_dtype) + ln_weight_cast = cast_if_needed(ln_weight, activation_dtype) + if ln_bias is not None: + ln_bias = cast_if_needed(ln_bias, activation_dtype) + if is_cpu_offload_enabled(): + start_offload(inputmat) + + tp_world_size = get_distributed_world_size(tp_group) + + # bwd needs fc1 input when grad is enabled, fc1 needs grad, and either + # 1) no checkpointing + # or 2) doing the recomputation with checkpointing + backwards_needs_fc1_input = fc1_weight_requires_grad and ( + (is_grad_enabled and not checkpoint) or is_recomputation + ) + + device = inp.device + + # Configure Userbuffers communication (comm+GEMM overlap) + ub_overlap_ag = ub_overlap_ag and is_grad_enabled and not return_layernorm_output_gathered + ub_overlap_rs = ub_overlap_rs and is_grad_enabled + + # Choose whether to use GEMM kernel with split accumulator + use_split_accumulator = _2X_ACC_FPROP + if fp8: + recipe = FP8GlobalStateManager.get_fp8_recipe() + if hasattr(recipe, "fp8_gemm_fprop"): + use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator + + # Configure quantizer for norm output + if fp8: + if fc1_input_quantizer is None: + raise ValueError("Missing quantizer for FC1 input tensor") + fc1_input_quantizer.set_usage(rowwise=True, columnwise=backwards_needs_fc1_input) + if sequence_parallel and fc1_input_quantizer.supports_only_rowwise_all_gather(): + # All-gather is not supported with FP8 column-wise data + fc1_input_quantizer.set_usage(columnwise=False) + # Amax reduction group for the FC1 input quantizer (column-parallel sequence parallel) + set_quantizer_amax_reduction_group( + fc1_input_quantizer, + tp_group if (sequence_parallel and set_parallel_mode) else None, + ) - # Choose whether to use GEMM kernel with split accumulator - use_split_accumulator = _2X_ACC_FPROP - if fp8: - recipe = FP8GlobalStateManager.get_fp8_recipe() - if hasattr(recipe, "fp8_gemm_fprop"): - use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator + # for fp8 DelayedScaling: layernorm output = FP8 + # only output of the linear is returned + # for return_layernorm_output: layernorm output = High precision, then cast to FP8 + # high precision layernorm output and output of the linear are returned + # 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 + ) + + # Apply normalization + ln_out, mu, rsigma = apply_normalization( + inputmat, + None, # ln_out + ln_weight_cast, + ln_bias, + args.eps, + fc1_input_quantizer if with_quantized_norm else None, + inputmat.dtype, + args.normalization, + args.fwd_ln_sm_margin, + args.zero_centered_gamma, + ) + ln_out_return = None + + # do not return layernorm output unless 1) no checkpointing or 2) checkpointing but not recomputing + if (return_layernorm_output or return_layernorm_output_gathered) and not is_recomputation: + ln_out_return = ln_out + + # Prepare GEMM input + # Note: Cast to expected dtype and perform tensor-parallel communication + ln_out_total = None + ub_obj_lnout = None + if sequence_parallel: + + # do not return ln output if checkpointing and in recomputation, not necessary + if return_layernorm_output_gathered and not is_recomputation: + # Perform all-gather in high precision if gathered + # norm output will be returned + ln_out_total, _ = gather_along_first_dim(ln_out, tp_group) + ln_out_return = ln_out_total + if fp8 or debug: + ln_out = fc1_input_quantizer(ln_out) + fc1_input_quantizer.set_usage(rowwise=True, columnwise=False) + ln_out_total = fc1_input_quantizer(ln_out_total) + else: + quantizer = None + if fp8 or debug: + quantizer = fc1_input_quantizer + # custom recipe doesn't need to support quantized AG + if not with_quantized_norm and not custom: + ln_out = fc1_input_quantizer(ln_out) + fc1_input_quantizer.set_usage(rowwise=True, columnwise=False) + if ub_overlap_ag: + # Copy into Userbuffers buffer + ub_obj_lnout = get_ub("fc1_fprop", fp8) + ln_out_total, _ = fill_userbuffers_buffer_for_all_gather( + ub_obj_lnout, + ln_out, + quantizer, + tp_group, + ) + else: + # All-gather with NCCL + ln_out_total, _ = gather_along_first_dim( + ln_out, + tp_group, + quantizer=quantizer, + ) + else: + if (fp8 or debug) and not with_quantized_norm: + ln_out = fc1_input_quantizer(ln_out) + ln_out_total = ln_out + + # Cast weights to expected dtype + new_fc1_weight_workspace = None + new_fc2_weight_workspace = None + fc1_weight_final = fc1_weight + fc2_weight_final = fc2_weight + # FSDP2: Skip columnwise/transpose creation during forward (not + # recompute) to avoid accumulating FP8 caches across layers. + # Backward's FSDP2 all-gather will recreate them. (Issue #2681) + fsdp2_skip_columnwise = is_fsdp2 and not is_recomputation + if fp8 or debug: + update_ws = is_first_microbatch is None or is_first_microbatch + # If weight is already quantized, weight._quantizer is its true quantizer. + # for debug mode we create quantizer every iteration, thus we need to set the quantizer states + if isinstance(fc1_weight, QuantizedTensorStorage) and not debug: + fc1_weight_quantizer = fc1_weight._quantizer + elif fc1_weight_quantizer is not None: + fc1_weight_quantizer.set_usage( + rowwise=True, + columnwise=is_grad_enabled and not fsdp2_skip_columnwise, + ) - # Configure quantizer for norm output - if fp8: - if fc1_input_quantizer is None: - raise ValueError("Missing quantizer for FC1 input tensor") - fc1_input_quantizer.set_usage(rowwise=True, columnwise=backwards_needs_fc1_input) - if sequence_parallel and fc1_input_quantizer.supports_only_rowwise_all_gather(): - # All-gather is not supported with FP8 column-wise data - fc1_input_quantizer.set_usage(columnwise=False) - # Amax reduction group for the FC1 input quantizer (column-parallel sequence parallel) - set_quantizer_amax_reduction_group( - fc1_input_quantizer, - tp_group if (sequence_parallel and set_parallel_mode) else None, + if isinstance(fc2_weight, QuantizedTensorStorage) and not debug: + fc2_weight_quantizer = fc2_weight._quantizer + elif fc2_weight_quantizer is not None: + fc2_weight_quantizer.set_usage( + rowwise=True, + columnwise=is_grad_enabled and not fsdp2_skip_columnwise, ) - # for fp8 DelayedScaling: layernorm output = FP8 - # only output of the linear is returned - # for return_layernorm_output: layernorm output = High precision, then cast to FP8 - # high precision layernorm output and output of the linear are returned - # 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 + fc1_weight_final, new_fc1_weight_workspace = quantize_weight( + tensor=fc1_weight, + quantizer=fc1_weight_quantizer, + workspace=args.fc1_weight_workspace, + update_workspace=update_ws, + skip_update_flag=args.skip_fp8_weight_update, + fsdp_group=fsdp_group, + workspace_dtype=activation_dtype, + cache=args.cache_weight, ) - - # Apply normalization - ln_out, mu, rsigma = apply_normalization( - inputmat, - None, # ln_out - ln_weight, - ln_bias, - eps, - fc1_input_quantizer if with_quantized_norm else None, - inputmat.dtype, - normalization, - fwd_ln_sm_margin, - zero_centered_gamma, + fc2_weight_final, new_fc2_weight_workspace = quantize_weight( + tensor=fc2_weight, + quantizer=fc2_weight_quantizer, + workspace=args.fc2_weight_workspace, + update_workspace=update_ws, + skip_update_flag=args.skip_fp8_weight_update, + fsdp_group=fsdp_group, + workspace_dtype=activation_dtype, + cache=args.cache_weight, ) - ln_out_return = None - - # do not return layernorm output unless 1) no checkpointing or 2) checkpointing but not recomputing - if (return_layernorm_output or return_layernorm_output_gathered) and not is_recomputation: - ln_out_return = ln_out - - # Prepare GEMM input - # Note: Cast to expected dtype and perform tensor-parallel communication - ln_out_total = None - ub_obj_lnout = None - if sequence_parallel: - - # do not return ln output if checkpointing and in recomputation, not necessary - if return_layernorm_output_gathered and not is_recomputation: - # Perform all-gather in high precision if gathered - # norm output will be returned - ln_out_total, _ = gather_along_first_dim(ln_out, tp_group) - ln_out_return = ln_out_total - if fp8 or debug: - ln_out = fc1_input_quantizer(ln_out) - fc1_input_quantizer.set_usage(rowwise=True, columnwise=False) - ln_out_total = fc1_input_quantizer(ln_out_total) + fc1_weight_final.update_usage(rowwise_usage=True) + fc2_weight_final.update_usage(rowwise_usage=True) + else: + fc1_weight_final = cast_if_needed(fc1_weight_final, activation_dtype) + fc2_weight_final = cast_if_needed(fc2_weight_final, activation_dtype) + + # Cast biases to expected dtype + bias_dtype = activation_dtype + if needs_quantized_gemm(ln_out_total) and activation_dtype == torch.float32: + # cuBLAS does not support FP8 GEMM with FP32 bias, so we cast to BF16 + bias_dtype = torch.bfloat16 + fc1_bias_cast = fc1_bias + fc2_bias_cast = fc2_bias + if fc1_bias is not None: + fc1_bias_cast = cast_if_needed(fc1_bias, bias_dtype) + if fc2_bias is not None: + fc2_bias_cast = cast_if_needed(fc2_bias, bias_dtype) + + # Calibrate quantizers if needed + if not fp8 and fp8_calibration: + if fc1_input_quantizer is not None: + fc1_input_quantizer.calibrate(ln_out_total) + if fc1_weight_quantizer is not None: + fc1_weight_quantizer.calibrate(fc1_weight) + + # ------------------------------------------------------ + # FC1 GEMM + # ------------------------------------------------------ + + # There are 2 fusions possible: + # - gemm_gelu_fusion - default for full precision, optional for fp8 - need to turn on gemm_gelu_fusion, + # - bias_gelu_fusion - only for full precision. + # If both gemm_gelu_fusion and bias_gelu_fusion are enabled, only bias_gelu_fusion will be performer + if activation != "gelu": + # blockwise scaled gemms don't support gemm_gelu_fusion in fwd. + gemm_gelu_fusion = bias_gelu_fusion = False + else: + if fp8: + assert not bias_gelu_fusion, "Bias gelu fusion is supported only for full precision" + else: + gemm_gelu_fusion = True + if gemm_gelu_fusion and bias_gelu_fusion: + gemm_gelu_fusion = False + if debug: + gemm_gelu_fusion = False + fc1_outputs = general_gemm( + fc1_weight_final, + ln_out_total, + quantization_params=( + fc2_input_quantizer + if gemm_gelu_fusion + else fc1_output_quantizer # fused gelu output is in fp8 + ), + out_dtype=activation_dtype, + bias=( + fc1_bias_cast if not bias_gelu_fusion else None + ), # otherwise bias is added later (fused with gelu) + gelu=gemm_gelu_fusion, + use_split_accumulator=use_split_accumulator, + ub=ub_obj_lnout, + ub_type=tex.CommOverlapType.AG if ub_overlap_ag else None, + ) + + # ------------------------------------------------------ + # Finished FC1 GEMM... + # ------------------------------------------------------ + + # Deallocate FC1 GEMM input tensor if no longer needed + # first part of if statement means that we only clear ln_out_total if + # 1) checkpointing and not recomputing (in the forward stage, not bwd recompute stage) + # 2) not checkpointing and grad disabled + # The `is not ln_out` guard avoids clearing the bwd-saved tensor when + # ln_out_total aliases ln_out (cuBLASMp AG-fprop path). + if ( + ((checkpoint and not is_recomputation) or not is_grad_enabled) + and ln_out_total is not ln_out_return + and ln_out_total is not ln_out + ): + clear_tensor_data(ln_out_total) + + # ACTIVATION - sometimes activation is fused with the GEMM above. + + fc1_out_without_bias = None + act_params = args.activation_params or {} + + if bias_gelu_fusion: + fc1_out = None + fc1_out_without_bias, *_ = fc1_outputs + act_out = bias_gelu_fused(fc1_out_without_bias, fc1_bias_cast) + elif gemm_gelu_fusion: + act_out, _, fc1_out, _ = fc1_outputs + elif debug: + fc1_out, *_ = fc1_outputs + act_out = activation_func(fc1_out, None, **act_params) + act_out = fc2_input_quantizer(act_out) + else: + fc1_out, *_ = fc1_outputs + if fp8: + recipe = FP8GlobalStateManager.get_fp8_recipe() + if recipe.float8_block_scaling(): + # tex.quantize does not support GELU fusion for blockwise + act_out = activation_func(fc1_out, None, **act_params) + act_out = tex.quantize(act_out, fc2_input_quantizer) + elif recipe.custom(): + # tex.quantize does not support custom quantizers + act_out = activation_func(fc1_out, None, **act_params) + act_out = fc2_input_quantizer(act_out) else: - quantizer = None - if fp8 or debug: - quantizer = fc1_input_quantizer - # custom recipe doesn't need to support quantized AG - if not with_quantized_norm and not custom: - ln_out = fc1_input_quantizer(ln_out) - fc1_input_quantizer.set_usage(rowwise=True, columnwise=False) - if ub_overlap_ag: - # Copy into Userbuffers buffer - ub_obj_lnout = get_ub("fc1_fprop", fp8) - ln_out_total, _ = fill_userbuffers_buffer_for_all_gather( - ub_obj_lnout, - ln_out, - quantizer, - tp_group, - ) - else: - # All-gather with NCCL - ln_out_total, _ = gather_along_first_dim( - ln_out, - tp_group, - quantizer=quantizer, - ) + act_out = activation_func(fc1_out, fc2_input_quantizer, **act_params) else: - if (fp8 or debug) and not with_quantized_norm: - ln_out = fc1_input_quantizer(ln_out) - ln_out_total = ln_out - - # Cast weights to expected dtype - new_fc1_weight_workspace = None - new_fc2_weight_workspace = None - fc1_weight_final = fc1_weight - fc2_weight_final = fc2_weight - # FSDP2: Skip columnwise/transpose creation during forward (not - # recompute) to avoid accumulating FP8 caches across layers. - # Backward's FSDP2 all-gather will recreate them. (Issue #2681) - fsdp2_skip_columnwise = is_fsdp2 and not is_recomputation - if fp8 or debug: - update_ws = is_first_microbatch is None or is_first_microbatch - # If weight is already quantized, weight._quantizer is its true quantizer. - # for debug mode we create quantizer every iteration, thus we need to set the quantizer states - if isinstance(fc1_weight, QuantizedTensorStorage) and not debug: - fc1_weight_quantizer = fc1_weight._quantizer - elif fc1_weight_quantizer is not None: - fc1_weight_quantizer.set_usage( - rowwise=True, - columnwise=is_grad_enabled and not fsdp2_skip_columnwise, - ) + if fp8_calibration: + act_out = activation_func(fc1_out, None, **act_params) + else: + act_out = activation_func(fc1_out, fc2_input_quantizer, **act_params) - if isinstance(fc2_weight, QuantizedTensorStorage) and not debug: - fc2_weight_quantizer = fc2_weight._quantizer - elif fc2_weight_quantizer is not None: - fc2_weight_quantizer.set_usage( - rowwise=True, - columnwise=is_grad_enabled and not fsdp2_skip_columnwise, - ) + if not fp8 and fp8_calibration: + if fc2_input_quantizer is not None: + fc2_input_quantizer.calibrate(act_out) + + # we want to skip fc2 computation if we are checkpointing and recomputing, + # otherwise we compute fc2 + fc2_out = None + if not (is_recomputation and checkpoint): + + # if we get to this point, we know this is not bwd recomputation + # so we must be in the fwd + # now is_grad_enabled can be true or false + # if false, can safely delete + # if true, we can only delete if checkpoint is true, since we will recompute anyways, + # otherwise, checkpoint is false, so cant delete + if checkpoint or not is_grad_enabled: # we can safely get rid of these if this is the case + clear_tensor_data(fc1_out) - fc1_weight_final, new_fc1_weight_workspace = quantize_weight( - tensor=fc1_weight, - quantizer=fc1_weight_quantizer, - workspace=fc1_weight_workspace, - update_workspace=update_ws, - skip_update_flag=skip_fp8_weight_update, - fsdp_group=fsdp_group, - workspace_dtype=activation_dtype, - cache=cache_weight, - ) - fc2_weight_final, new_fc2_weight_workspace = quantize_weight( - tensor=fc2_weight, - quantizer=fc2_weight_quantizer, - workspace=fc2_weight_workspace, - update_workspace=update_ws, - skip_update_flag=skip_fp8_weight_update, - fsdp_group=fsdp_group, - workspace_dtype=activation_dtype, - cache=cache_weight, - ) - fc1_weight_final.update_usage(rowwise_usage=True) - fc2_weight_final.update_usage(rowwise_usage=True) - else: - fc1_weight_final = cast_if_needed(fc1_weight_final, activation_dtype) - fc2_weight_final = cast_if_needed(fc2_weight_final, activation_dtype) - - # Cast biases to expected dtype - bias_dtype = activation_dtype - if needs_quantized_gemm(ln_out_total) and activation_dtype == torch.float32: - # cuBLAS does not support FP8 GEMM with FP32 bias, so we cast to BF16 - bias_dtype = torch.bfloat16 - if fc1_bias is not None: - fc1_bias = cast_if_needed(fc1_bias, bias_dtype) - if fc2_bias is not None: - fc2_bias = cast_if_needed(fc2_bias, bias_dtype) - - # Calibrate quantizers if needed if not fp8 and fp8_calibration: - if fc1_input_quantizer is not None: - fc1_input_quantizer.calibrate(ln_out_total) - if fc1_weight_quantizer is not None: - fc1_weight_quantizer.calibrate(fc1_weight) + + if fc2_weight_quantizer is not None: + fc2_weight_quantizer.calibrate(fc2_weight) + + # Configure Userbuffers reduce-scatter if needed + ub_obj_fc2out = None + reduce_scatter_out = None + if ub_overlap_rs: + ub_obj_fc2out = get_ub("fc2_fprop", fp8) + dim_size = list(act_out.size()) + dim_size[0] //= tp_world_size + dim_size[-1] = fc2_weight.size(0) + reduce_scatter_out = torch.empty(dim_size, dtype=activation_dtype, device=device) # ------------------------------------------------------ - # FC1 GEMM + # FC2 GEMM # ------------------------------------------------------ - - # There are 2 fusions possible: - # - gemm_gelu_fusion - default for full precision, optional for fp8 - need to turn on gemm_gelu_fusion, - # - bias_gelu_fusion - only for full precision. - # If both gemm_gelu_fusion and bias_gelu_fusion are enabled, only bias_gelu_fusion will be performer - if activation != "gelu": - # blockwise scaled gemms don't support gemm_gelu_fusion in fwd. - gemm_gelu_fusion = bias_gelu_fusion = False - else: - if fp8: - assert not bias_gelu_fusion, "Bias gelu fusion is supported only for full precision" - else: - gemm_gelu_fusion = True - if gemm_gelu_fusion and bias_gelu_fusion: - gemm_gelu_fusion = False - if debug: - gemm_gelu_fusion = False - fc1_outputs = general_gemm( - fc1_weight_final, - ln_out_total, - quantization_params=( - fc2_input_quantizer - if gemm_gelu_fusion - else fc1_output_quantizer # fused gelu output is in fp8 - ), + gemm_out, *_, reduce_scatter_out = general_gemm( + fc2_weight_final, + act_out, out_dtype=activation_dtype, - bias=( - fc1_bias if not bias_gelu_fusion else None - ), # otherwise bias is added later (fused with gelu) - gelu=gemm_gelu_fusion, + bias=fc2_bias_cast, + quantization_params=fc2_output_quantizer, use_split_accumulator=use_split_accumulator, - ub=ub_obj_lnout, - ub_type=tex.CommOverlapType.AG if ub_overlap_ag else None, + ub=ub_obj_fc2out, + ub_type=tex.CommOverlapType.RS if ub_overlap_rs else None, + extra_output=reduce_scatter_out, ) - # ------------------------------------------------------ - # Finished FC1 GEMM... + # Finished FC2 GEMM... # ------------------------------------------------------ - # Deallocate FC1 GEMM input tensor if no longer needed - # first part of if statement means that we only clear ln_out_total if - # 1) checkpointing and not recomputing (in the forward stage, not bwd recompute stage) - # 2) not checkpointing and grad disabled - # The `is not ln_out` guard avoids clearing the bwd-saved tensor when - # ln_out_total aliases ln_out (cuBLASMp AG-fprop path). - if ( - ((checkpoint and not is_recomputation) or not is_grad_enabled) - and ln_out_total is not ln_out_return - and ln_out_total is not ln_out - ): - clear_tensor_data(ln_out_total) - - # ACTIVATION - sometimes activation is fused with the GEMM above. - - fc1_out_without_bias = None - act_params = activation_params or {} - - if bias_gelu_fusion: - fc1_out = None - fc1_out_without_bias, *_ = fc1_outputs - act_out = bias_gelu_fused(fc1_out_without_bias, fc1_bias) - elif gemm_gelu_fusion: - act_out, _, fc1_out, _ = fc1_outputs - elif debug: - fc1_out, *_ = fc1_outputs - act_out = activation_func(fc1_out, None, **act_params) - act_out = fc2_input_quantizer(act_out) - else: - fc1_out, *_ = fc1_outputs - if fp8: - recipe = FP8GlobalStateManager.get_fp8_recipe() - if recipe.float8_block_scaling(): - # tex.quantize does not support GELU fusion for blockwise - act_out = activation_func(fc1_out, None, **act_params) - act_out = tex.quantize(act_out, fc2_input_quantizer) - elif recipe.custom(): - # tex.quantize does not support custom quantizers - act_out = activation_func(fc1_out, None, **act_params) - act_out = fc2_input_quantizer(act_out) - else: - act_out = activation_func(fc1_out, fc2_input_quantizer, **act_params) - else: - if fp8_calibration: - act_out = activation_func(fc1_out, None, **act_params) - else: - act_out = activation_func(fc1_out, fc2_input_quantizer, **act_params) - - if not fp8 and fp8_calibration: - if fc2_input_quantizer is not None: - fc2_input_quantizer.calibrate(act_out) - - # we want to skip fc2 computation if we are checkpointing and recomputing, - # otherwise we compute fc2 - if not (is_recomputation and checkpoint): - - # if we get to this point, we know this is not bwd recomputation - # so we must be in the fwd - # now is_grad_enabled can be true or false - # if false, can safely delete - # if true, we can only delete if checkpoint is true, since we will recompute anyways, - # otherwise, checkpoint is false, so cant delete - if ( - checkpoint or not is_grad_enabled - ): # we can safely get rid of these if this is the case - clear_tensor_data(fc1_out) - - if not fp8 and fp8_calibration: - - if fc2_weight_quantizer is not None: - fc2_weight_quantizer.calibrate(fc2_weight) - - # Configure Userbuffers reduce-scatter if needed - ub_obj_fc2out = None - reduce_scatter_out = None - if ub_overlap_rs: - ub_obj_fc2out = get_ub("fc2_fprop", fp8) - dim_size = list(act_out.size()) - dim_size[0] //= tp_world_size - dim_size[-1] = fc2_weight.size(0) - reduce_scatter_out = torch.empty(dim_size, dtype=activation_dtype, device=device) - - # ------------------------------------------------------ - # FC2 GEMM - # ------------------------------------------------------ - gemm_out, *_, reduce_scatter_out = general_gemm( - fc2_weight_final, - act_out, - out_dtype=activation_dtype, - bias=fc2_bias, - quantization_params=fc2_output_quantizer, - use_split_accumulator=use_split_accumulator, - ub=ub_obj_fc2out, - ub_type=tex.CommOverlapType.RS if ub_overlap_rs else None, - extra_output=reduce_scatter_out, + # Deallocate tensors if no longer needed, again, can safely deallocate + if checkpoint or not is_grad_enabled: # same logic as last clear_tensor_data block + clear_tensor_data(act_out, fc1_out_without_bias, fc1_out) + + # Prepare output tensor + # Note: Perform tensor-parallel communication if needed + if ub_overlap_rs: + # cuBLASMp writes the reduce-scattered output directly into the + # GEMM output tensor; Userbuffers writes it into the extra-output buffer. + fc2_out = ( + gemm_out + if ub_obj_fc2out is not None and ub_obj_fc2out.with_cublasmp() + else reduce_scatter_out ) - # ------------------------------------------------------ - # Finished FC2 GEMM... - # ------------------------------------------------------ - - # Deallocate tensors if no longer needed, again, can safely deallocate - if checkpoint or not is_grad_enabled: # same logic as last clear_tensor_data block - clear_tensor_data(act_out, fc1_out_without_bias, fc1_out) - - # Prepare output tensor - # Note: Perform tensor-parallel communication if needed - fc2_out = None - if ub_overlap_rs: - # cuBLASMp writes the reduce-scattered output directly into the - # GEMM output tensor; Userbuffers writes it into the extra-output buffer. - fc2_out = ( - gemm_out - if ub_obj_fc2out is not None and ub_obj_fc2out.with_cublasmp() - else reduce_scatter_out + elif set_parallel_mode and sequence_parallel: + fc2_out, _ = reduce_scatter_along_first_dim(gemm_out, tp_group) + elif set_parallel_mode and tensor_parallel: + if args.symmetric_ar_type is not None: + fc2_out, _ = symmetric_all_reduce( + gemm_out, tp_group, all_reduce_type=args.symmetric_ar_type ) - elif set_parallel_mode and sequence_parallel: - fc2_out, _ = reduce_scatter_along_first_dim(gemm_out, tp_group) - elif set_parallel_mode and tensor_parallel: - if symmetric_ar_type is not None: - fc2_out, _ = symmetric_all_reduce( - gemm_out, tp_group, all_reduce_type=symmetric_ar_type - ) - else: - fc2_out, _ = allreduce(gemm_out, tp_group) else: - fc2_out = gemm_out - fc2_out = fc2_out.view(-1, *inp_shape[1:-1], fc2_out.shape[-1]) - - # now saving stuff for bwd: - # if we are using checkpointing, this information will be saved in the bwd recomputation stage, so can skip it in fwd - # if we are not checkpointing, then we must save this if grad is enabled - if is_grad_enabled and not save_for_checkpoint: - - ctx.fc1_weight_quantizer = fc1_weight_quantizer - ctx.fc2_weight_quantizer = fc2_weight_quantizer - - if not fc1_weight.requires_grad: - if not return_layernorm_output: - clear_tensor_data(ln_out) - ln_out = None - if not fc2_weight.requires_grad: - clear_tensor_data(act_out) - act_out = None - - if not checkpoint: # regular path, no selective activation checkpointing - - if cpu_offloading: - mark_activation_offload( - inputmat, mu, rsigma, ln_out, fc1_out, fc1_out_without_bias, act_out - ) - - # Scatter intermediate/activation tensors saved for the backward pass - # NOTE: weight_fp8 = weight when ctx.fp8 == False and torch.disttributed.FSDP already - # shards/unshards the base weights so we don't do it ourselves - ctx.fsdp_group = fsdp_group - - ctx.fsdp_shapes = ( - _fsdp_scatter_tensors( # again, ony relevant if we have activations to save - fsdp_group, - mu, - rsigma, - ln_out, - fc1_out_without_bias if bias_gelu_fusion else fc1_out, - act_out, - ( - fc1_weight_final - if fp8 and not isinstance(fc1_weight, Float8Tensor) - else None - ), - ( - fc2_weight_final - if fp8 and not isinstance(fc2_weight, Float8Tensor) - else None - ), - ) + fc2_out, _ = allreduce(gemm_out, tp_group) + else: + fc2_out = gemm_out + fc2_out = fc2_out.view(-1, *inp_shape[1:-1], fc2_out.shape[-1]) + + # now saving stuff for bwd: + # if we are using checkpointing, this information will be saved in the bwd recomputation stage, so can skip it in fwd + # if we are not checkpointing, then we must save this if grad is enabled + if is_grad_enabled and not save_for_checkpoint: + if ctx_attrs is None: + ctx_attrs = {} + + if not fc1_weight_requires_grad: + if not return_layernorm_output: + clear_tensor_data(ln_out) + ln_out = None + if not fc2_weight_requires_grad: + clear_tensor_data(act_out) + act_out = None + + fsdp_shapes = None + if not checkpoint: # regular path, no selective activation checkpointing + + if cpu_offloading: + mark_activation_offload( + inputmat, mu, rsigma, ln_out, fc1_out, fc1_out_without_bias, act_out ) - if cpu_offloading: - mark_not_offload( - ln_weight, - ln_bias, - fc1_weight_final, - fc1_weight, - fc1_bias, - fc2_weight_final, - fc2_weight, - fc2_bias, - ) - # FSDP2: Don't save FP8 workspace copies for non-quantized - # weights. Backward will re-quantize from the FSDP2 - # all-gathered weight parameter. (Issue #2681) - fc1_wt_save = fc1_weight_final - fc2_wt_save = fc2_weight_final - if fsdp2_skip_columnwise: - if fc1_weight_final is not fc1_weight: - fc1_wt_save = None - if fc2_weight_final is not fc2_weight: - fc2_wt_save = None - tensors_to_save, tensor_objects = prepare_for_saving( - inputmat, - ln_weight, - ln_out, - fc1_wt_save, - fc1_weight, - fc1_bias, - fc1_out, - fc1_out_without_bias, - act_out, - fc2_wt_save, - fc2_weight, - fc2_bias, + # Scatter intermediate/activation tensors saved for the backward pass + # NOTE: weight_fp8 = weight when fp8 == False and torch.disttributed.FSDP already + # shards/unshards the base weights so we don't do it ourselves + fsdp_shapes = ( + _fsdp_scatter_tensors( # again, ony relevant if we have activations to save + fsdp_group, mu, rsigma, + ln_out, + fc1_out_without_bias if bias_gelu_fusion else fc1_out, + act_out, + ( + fc1_weight_final + if fp8 and not isinstance(fc1_weight, Float8Tensor) + else None + ), + ( + fc2_weight_final + if fp8 and not isinstance(fc2_weight, Float8Tensor) + else None + ), ) - - ctx.save_for_backward(*tensors_to_save) - ctx.tensor_objects = tensor_objects - - if fuse_wgrad_accumulation: - # Keep weakrefs to weights to preserve attributes like main_grad - # when we need to modify the weight python objects - ctx.fc1_weight_python_object_ref = ( - weakref.ref(fc1_weight) if fc1_weight.requires_grad else None - ) - ctx.fc2_weight_python_object_ref = ( - weakref.ref(fc2_weight) if fc2_weight.requires_grad else None - ) - ctx.fc1_weight_overwrites_main_grad = getattr( - fc1_weight, "overwrite_main_grad", False - ) - ctx.fc2_weight_overwrites_main_grad = getattr( - fc2_weight, "overwrite_main_grad", False - ) - # This check is needed to ensure that main_grad is not created - # during the forward pass when using MCore FSDP as it creates - # the main_grad buffer lazily before backprop - if hasattr(fc1_weight, "__fsdp_param__") and hasattr(fc2_weight, "__fsdp_param__"): - # MCore FSDP creates main_grad lazily before backward - ctx.fc1_main_grad_func = ( - fc1_weight.get_main_grad if fc1_weight.requires_grad else lambda: None - ) - ctx.fc2_main_grad_func = ( - fc2_weight.get_main_grad if fc2_weight.requires_grad else lambda: None - ) - else: - ctx.fc1_main_grad_func = lambda: fc1_weight.main_grad - ctx.fc2_main_grad_func = lambda: fc2_weight.main_grad - - ctx.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None - ctx.backward_override = backward_override - ctx.fc1_grad_input_quantizer = fc1_grad_input_quantizer - ctx.fc1_grad_weight_quantizer = fc1_grad_weight_quantizer - ctx.fc1_grad_output_quantizer = fc1_grad_output_quantizer - ctx.fc2_grad_input_quantizer = fc2_grad_input_quantizer - ctx.fc2_grad_weight_quantizer = fc2_grad_weight_quantizer - ctx.fc2_grad_output_quantizer = fc2_grad_output_quantizer - ctx.fc1_input_quantizer = fc1_input_quantizer - ctx.fc2_input_quantizer = fc2_input_quantizer - - ctx.fc1_weight_requires_grad = fc1_weight.requires_grad - ctx.fc2_weight_requires_grad = fc2_weight.requires_grad - ctx.fc1_weight = fc1_weight - ctx.fc2_weight = fc2_weight - ctx.fsdp2_skip_columnwise = fsdp2_skip_columnwise - # Store raw is_fsdp2 flag for backward cleanup — must not be - # gated on is_recomputation since backward cleanup runs after - # the real backward, not the recomputation forward. - ctx.is_fsdp2 = is_fsdp2 - - ctx.device = device - ctx.activation_dtype = activation_dtype - ctx.activation = activation - ctx.activation_params = activation_params - ctx.fp8 = fp8 - ctx.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None - ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation - ctx.cpu_offloading = cpu_offloading - ctx.is_first_microbatch = is_first_microbatch - ctx.use_bias = fc2_bias is not None - ctx.sequence_parallel = sequence_parallel - ctx.tensor_parallel = tensor_parallel - ctx.inp_shape = inp_shape - ctx.tp_group = tp_group - ctx.tp_size = tp_size - ctx.bias_gelu_fusion = bias_gelu_fusion - ctx.return_layernorm_output = return_layernorm_output - ctx.return_layernorm_output_gathered = ( - return_layernorm_output_gathered and sequence_parallel - ) - ctx.set_parallel_mode = set_parallel_mode - ctx.bwd_ln_sm_margin = bwd_ln_sm_margin - ctx.zero_centered_gamma = zero_centered_gamma - ctx.ub_bulk_wgrad = ub_bulk_wgrad - ctx.ub_bulk_dgrad = ub_bulk_dgrad - ctx.ub_overlap_rs_dgrad = ub_overlap_rs_dgrad - ctx.ub_overlap_ag = ub_overlap_ag - ctx.debug = debug - - ctx.requires_dgrad = ( - inp.requires_grad or ln_weight.requires_grad or ln_bias.requires_grad ) - ctx.normalization = normalization - ctx.reduce_and_update_bwd_fp8_tensors = False - if ctx.fp8 and requires_grad( - inp, ln_weight, ln_bias, fc1_weight, fc2_weight, fc1_bias, fc2_bias - ): - qstate = FP8GlobalStateManager.quantization_state - _first_fp8_module = qstate.is_first_fp8_module - ctx.reduce_and_update_bwd_fp8_tensors = FP8GlobalStateManager.is_first_fp8_module() - if in_fp8_activation_recompute_phase() or is_recomputation: - qstate.is_first_fp8_module = _first_fp8_module - - ctx.wgrad_store = wgrad_store - if is_recomputation: # return the recomputed tensors - return ( - ctx, - inputmat, - ln_weight, - ln_out, + + if cpu_offloading: + mark_not_offload( + ln_weight_cast, + ln_bias, fc1_weight_final, fc1_weight, - fc1_bias, - fc1_out, - fc1_out_without_bias, - act_out, + fc1_bias_cast, fc2_weight_final, fc2_weight, - fc2_bias, - mu, - rsigma, + fc2_bias_cast, ) - # we only get to this point if we are not recomputing for bwd, since that would have returned in the block above - ln_out_for_return = None - if return_layernorm_output: - if return_layernorm_output_gathered: - shape = list(inp_shape) - shape[0] *= tp_size if (sequence_parallel and set_parallel_mode) else 1 - ln_out_for_return = ln_out_return.view(shape) - else: - ln_out_for_return = ln_out_return.view(inp_shape) - return fc2_out, ln_out_for_return, new_fc1_weight_workspace, new_fc2_weight_workspace - - @staticmethod - def forward( - ctx, - inp: torch.Tensor, - ln_weight: torch.Tensor, - ln_bias: torch.Tensor, - fc1_weight: torch.Tensor, - fc1_weight_workspace: Optional[torch.Tensor], - fc1_bias: torch.Tensor, - fc2_weight: torch.Tensor, - fc2_weight_workspace: Optional[torch.Tensor], - fc2_bias: torch.Tensor, - non_tensor_args: Tuple, - ) -> Tuple[torch.Tensor, ...]: - # pylint: disable=missing-function-docstring - - # add recompute_for_bwd - non_tensor_args += (False,) - - return _LayerNormMLP._forward( - ctx, - inp, - ln_weight, - ln_bias, + # FSDP2: Don't save FP8 workspace copies for non-quantized + # weights. Backward will re-quantize from the FSDP2 + # all-gathered weight parameter. (Issue #2681) + fc1_wt_save = fc1_weight_final + fc2_wt_save = fc2_weight_final + if fsdp2_skip_columnwise: + if fc1_weight_final is not fc1_weight: + fc1_wt_save = None + if fc2_weight_final is not fc2_weight: + fc2_wt_save = None + + # Dedup save slots that alias forward inputs or other outputs; + # ``_layernorm_mlp_setup_ctx`` rebuilds the refs. + def _wt_alias(wt_save, weight, new_workspace, workspace, prefix): + if wt_save is None: + return None + if wt_save is weight: + return prefix + "_weight" + if new_workspace is not None and wt_save is new_workspace: + return "new_" + prefix + "_weight_workspace" + if workspace is not None and wt_save is workspace: + return prefix + "_weight_workspace" + return None + + saved_tensor_aliases = ( + "inp" if inputmat is inp else None, + "ln_weight" if ln_weight_cast is ln_weight else None, + ( + "ln_out" + if return_layernorm_output + and ln_out is not None + and ln_out_return is not None + and ln_out is ln_out_return + else None + ), + _wt_alias( + fc1_wt_save, + fc1_weight, + new_fc1_weight_workspace, + args.fc1_weight_workspace, + "fc1", + ), + "fc1_weight", + "fc1_bias" if fc1_bias_cast is not None and fc1_bias_cast is fc1_bias else None, + None, + None, + None, + _wt_alias( + fc2_wt_save, + fc2_weight, + new_fc2_weight_workspace, + args.fc2_weight_workspace, + "fc2", + ), + "fc2_weight", + "fc2_bias" if fc2_bias_cast is not None and fc2_bias_cast is fc2_bias else None, + None, + None, + ) + saved = ( + inputmat, + ln_weight_cast, + ln_out, + fc1_wt_save, fc1_weight, - fc1_weight_workspace, - fc1_bias, + fc1_bias_cast, + fc1_out, + fc1_out_without_bias, + act_out, + fc2_wt_save, fc2_weight, - fc2_weight_workspace, - fc2_bias, - non_tensor_args, + fc2_bias_cast, + mu, + rsigma, + ) + tensors_to_save_from_forward = tuple( + None if alias is not None else tensor + for alias, tensor in zip(saved_tensor_aliases, saved) + ) + ctx_attrs["saved_tensor_aliases"] = saved_tensor_aliases + ctx_attrs["fsdp_shapes"] = fsdp_shapes + ctx_attrs["is_recomputation"] = is_recomputation + + if is_recomputation: # return the recomputed tensors + return None, None, None, None, tensors_to_save_from_forward, ctx_attrs + + # we only get to this point if we are not recomputing for bwd, since that would have returned in the block above + ln_out_for_return = None + if return_layernorm_output: + if return_layernorm_output_gathered: + shape = list(inp_shape) + shape[0] *= tp_size if (sequence_parallel and set_parallel_mode) else 1 + ln_out_for_return = ln_out_return.view(shape) + else: + ln_out_for_return = ln_out_return.view(inp_shape) + return ( + fc2_out, + ln_out_for_return, + new_fc1_weight_workspace, + new_fc2_weight_workspace, + tensors_to_save_from_forward, + ctx_attrs, + ) + + +def _layernorm_mlp_setup_ctx( + bwd_args: LayerNormMLPBwdArgs, + fwd_args: LayerNormMLPFwdArgs, + fwd_outputs: Tuple[Any, ...], + ctx_attrs: Dict, + tensors_to_save_from_forward: Tuple[Any, ...], +) -> Tuple[Any, ...]: + """Populate ``bwd_args`` from forward state. + + Returns the tensors that should be passed through ``prepare_for_saving`` + by the caller. ``fwd_outputs`` is ``(fc2_out, ln_out_return, + new_fc1_weight_workspace, new_fc2_weight_workspace)``; the last three + rebuild the deduped save slots. + + With ``ctx_attrs["checkpoint"]`` and no recomputation yet, only the + forward inputs are saved; the rest of ``bwd_args`` is filled when the + backward recomputes the forward. + """ + checkpoint = ctx_attrs.get("checkpoint", False) + is_recomputation = ctx_attrs.get("is_recomputation", False) + if checkpoint and not is_recomputation: + bwd_args.checkpoint = True + bwd_args.autocast_state = ctx_attrs["autocast_state"] + bwd_args.cpu_rng_state = ctx_attrs["cpu_rng_state"] + bwd_args.cuda_rng_state = ctx_attrs["cuda_rng_state"] + # The saved inputs are re-bound to the args when the backward recomputes. + bwd_args.checkpoint_fwd_args = dataclass_replace( + fwd_args, + inp=None, + ln_weight=None, + ln_bias=None, + fc1_weight=None, + fc1_bias=None, + fc2_weight=None, + fc2_bias=None, + fc1_weight_workspace=None, + fc2_weight_workspace=None, + cache_weight=False, + recompute_for_bwd=True, + ) + return tuple(getattr(fwd_args, name) for name in ctx_attrs["saved_tensor_aliases"]) + + inp = fwd_args.inp + fc1_weight = fwd_args.fc1_weight + fc2_weight = fwd_args.fc2_weight + fp8 = fwd_args.fp8 + debug = fwd_args.debug + fuse_wgrad_accumulation = fwd_args.fuse_wgrad_accumulation + fc1_weight_requires_grad = fwd_args.fc1_weight_requires_grad + fc2_weight_requires_grad = fwd_args.fc2_weight_requires_grad + + # Quantizers + bwd_args.fc1_input_quantizer = fwd_args.fc1_input_quantizer + bwd_args.fc2_input_quantizer = fwd_args.fc2_input_quantizer + bwd_args.fc1_weight_quantizer = ( + fc1_weight._quantizer + if (fp8 or debug) and isinstance(fc1_weight, QuantizedTensorStorage) and not debug + else fwd_args.fc1_weight_quantizer + ) + bwd_args.fc2_weight_quantizer = ( + fc2_weight._quantizer + if (fp8 or debug) and isinstance(fc2_weight, QuantizedTensorStorage) and not debug + else fwd_args.fc2_weight_quantizer + ) + bwd_args.fc1_grad_input_quantizer = fwd_args.fc1_grad_input_quantizer + bwd_args.fc1_grad_weight_quantizer = fwd_args.fc1_grad_weight_quantizer + bwd_args.fc1_grad_output_quantizer = fwd_args.fc1_grad_output_quantizer + bwd_args.fc2_grad_input_quantizer = fwd_args.fc2_grad_input_quantizer + bwd_args.fc2_grad_weight_quantizer = fwd_args.fc2_grad_weight_quantizer + bwd_args.fc2_grad_output_quantizer = fwd_args.fc2_grad_output_quantizer + + # Differentiability summary + bwd_args.use_bias = fwd_args.fc2_bias is not None + bwd_args.requires_dgrad = ( + fwd_args.input_requires_grad + or fwd_args.ln_weight_requires_grad + or fwd_args.ln_bias_requires_grad + ) + bwd_args.fc1_weight_requires_grad = fc1_weight_requires_grad + bwd_args.fc1_bias_requires_grad = fwd_args.fc1_bias_requires_grad + bwd_args.fc2_weight_requires_grad = fc2_weight_requires_grad + bwd_args.inp_shape = inp.shape + + # Normalization + bwd_args.normalization = fwd_args.normalization + bwd_args.zero_centered_gamma = fwd_args.zero_centered_gamma + bwd_args.bwd_ln_sm_margin = fwd_args.bwd_ln_sm_margin + bwd_args.return_layernorm_output = fwd_args.return_layernorm_output + bwd_args.return_layernorm_output_gathered = ( + fwd_args.return_layernorm_output_gathered and fwd_args.sequence_parallel + ) + + # Activation + bwd_args.activation = fwd_args.activation + bwd_args.activation_params = fwd_args.activation_params + bwd_args.bias_gelu_fusion = fwd_args.bias_gelu_fusion + + # Numerical / dtype config + bwd_args.activation_dtype = fwd_args.activation_dtype + bwd_args.fp8 = fp8 + bwd_args.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None + bwd_args.dgrad_use_split_accumulator = fwd_args.dgrad_use_split_accumulator + bwd_args.wgrad_use_split_accumulator = fwd_args.wgrad_use_split_accumulator + bwd_args.backward_override = fwd_args.backward_override + bwd_args.debug = debug + + # Tensor / sequence parallelism + bwd_args.set_parallel_mode = fwd_args.set_parallel_mode + bwd_args.tp_group = fwd_args.tp_group + bwd_args.tp_size = fwd_args.tp_size + bwd_args.tensor_parallel = fwd_args.tensor_parallel + bwd_args.sequence_parallel = fwd_args.sequence_parallel + + # Userbuffers + bwd_args.ub_overlap_ag = ( + fwd_args.ub_overlap_ag + and fwd_args.is_grad_enabled + and not fwd_args.return_layernorm_output_gathered + ) + bwd_args.ub_overlap_rs_dgrad = fwd_args.ub_overlap_rs_dgrad + bwd_args.ub_bulk_dgrad = fwd_args.ub_bulk_dgrad + bwd_args.ub_bulk_wgrad = fwd_args.ub_bulk_wgrad + + # FSDP + bwd_args.fsdp_group = fwd_args.fsdp_group + bwd_args.fsdp_shapes = ctx_attrs["fsdp_shapes"] + bwd_args.is_fsdp2 = fwd_args.is_fsdp2 + + # Weight-grad scheduling / accumulation + bwd_args.is_first_microbatch = fwd_args.is_first_microbatch + bwd_args.fuse_wgrad_accumulation = fuse_wgrad_accumulation + bwd_args.wgrad_store = fwd_args.wgrad_store + if fuse_wgrad_accumulation: + # Keep weakrefs to weights to preserve attributes like main_grad + # when we need to modify the weight python objects + bwd_args.fc1_weight_ref = weakref.ref(fc1_weight) if fc1_weight_requires_grad else None + bwd_args.fc2_weight_ref = weakref.ref(fc2_weight) if fc2_weight_requires_grad else None + bwd_args.fc1_weight_overwrites_main_grad = getattr(fc1_weight, "overwrite_main_grad", False) + bwd_args.fc2_weight_overwrites_main_grad = getattr(fc2_weight, "overwrite_main_grad", False) + # MCore FSDP creates main_grad lazily before backward, so don't touch it here + if hasattr(fc1_weight, "__fsdp_param__") and hasattr(fc2_weight, "__fsdp_param__"): + bwd_args.fc1_main_grad_func = ( + fc1_weight.get_main_grad if fc1_weight_requires_grad else lambda: None + ) + bwd_args.fc2_main_grad_func = ( + fc2_weight.get_main_grad if fc2_weight_requires_grad else lambda: None + ) + else: + bwd_args.fc1_main_grad_func = lambda: fc1_weight.main_grad + bwd_args.fc2_main_grad_func = lambda: fc2_weight.main_grad + + # Misc + bwd_args.cpu_offloading = fwd_args.cpu_offloading + + saved = list(tensors_to_save_from_forward) + aliases = ctx_attrs["saved_tensor_aliases"] + in_features = inp.shape[-1] + for i, alias in enumerate(aliases): + if alias is None: + continue + if alias == "inp": + saved[i] = inp.view((-1, in_features)) + elif alias == "ln_out": + saved[i] = fwd_outputs[1].view((-1, in_features)) + elif alias == "new_fc1_weight_workspace": + saved[i] = fwd_outputs[2] + elif alias == "new_fc2_weight_workspace": + saved[i] = fwd_outputs[3] + else: + saved[i] = getattr(fwd_args, alias) + if fwd_args.cpu_offloading: + # Rebuilt views don't carry the offload marks set on the forward tensors + mark_activation_offload( + *(saved[i] for i, alias in enumerate(aliases) if alias in ("inp", "ln_out")) + ) + return tuple(saved) + + +def _layernorm_mlp_recompute( + bwd_args: LayerNormMLPBwdArgs, ctx: torch.autograd.function.FunctionCtx +) -> None: + """Bind the saved tensors to ``bwd_args``, recomputing the forward when + the module ran with activation checkpointing.""" + tensors = restore_from_func_ctx(ctx) + if not bwd_args.checkpoint: + bwd_args.set_saved_tensors(tensors) + return + + # backward is not in autocast context, so we set the state here + # we also have to set the quantizer states to what they were before the forward pass (only relevant for DelayedScaling recipe) + fwd_args: LayerNormMLPFwdArgs = bwd_args.checkpoint_fwd_args + bwd_args.checkpoint_fwd_args = None + final_autocast_state = FP8GlobalStateManager.get_autocast_state() + FP8GlobalStateManager.set_autocast_state(bwd_args.autocast_state) + if ( + fwd_args.fp8 + and FP8GlobalStateManager.get_fp8_recipe().__class__.__name__ == "DelayedScaling" + ): # only applicable for delayed scaling + FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute( + fwd_args.fp8_meta + ) # set old quantizer state + + # get current rng state + final_cpu_rng_state = torch.get_rng_state() + final_cuda_rng_state = _get_cuda_rng_state() + + # set rng state for fwd + torch.set_rng_state(bwd_args.cpu_rng_state) + _set_cuda_rng_state(bwd_args.cuda_rng_state) + + for name, tensor in zip(_CHECKPOINT_SAVED_ALIASES, tensors): + setattr(fwd_args, name, tensor) + ( + _, + _, + new_fc1_weight_workspace, + new_fc2_weight_workspace, + tensors_to_save_from_forward, + ctx_attrs, + ) = _layernorm_mlp_forward_impl(fwd_args) + recomputed = _layernorm_mlp_setup_ctx( + bwd_args, + fwd_args, + (None, None, new_fc1_weight_workspace, new_fc2_weight_workspace), + ctx_attrs, + tensors_to_save_from_forward, + ) + bwd_args.set_saved_tensors(recomputed) + if fwd_args.fp8 and fwd_args.any_requires_grad(): + bwd_args.reduce_and_update_bwd_fp8_tensors = check_fp8_reduce_and_update( + restore_first_module=True ) - @staticmethod - def _recompute(ctx): - # pylint: disable=missing-function-docstring - - tensors = restore_from_func_ctx(ctx) + FP8GlobalStateManager.set_autocast_state(final_autocast_state) + if ( + fwd_args.fp8 + and FP8GlobalStateManager.get_fp8_recipe().__class__.__name__ == "DelayedScaling" + ): + FP8GlobalStateManager.restore_fp8_meta_tensors(fwd_args.fp8_meta) # restore quantizers - if ctx.checkpoint: # do recomputation from the original args + # set rng state for fwd + torch.set_rng_state(final_cpu_rng_state) + _set_cuda_rng_state(final_cuda_rng_state) - # backward is not in autocast context, so we set the state here - # we also have to set the quantizer states to what they were before the forward pass (only relevant for DelayedScaling recipe) - final_autocast_state = FP8GlobalStateManager.get_autocast_state() - FP8GlobalStateManager.set_autocast_state(ctx.autocast_state) - if ( - ctx.other_args["fp8"] - and FP8GlobalStateManager.get_fp8_recipe().__class__.__name__ == "DelayedScaling" - ): # only applicable for delayed scaling - FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute( - ctx.other_args["fp8_meta"] - ) # set old quantizer state - # get current rng state - final_cpu_rng_state = torch.get_rng_state() - final_cuda_rng_state = _get_cuda_rng_state() +def _layernorm_mlp_backward_impl( + args: LayerNormMLPBwdArgs, +) -> Tuple[Union[torch.Tensor, None], ...]: + """Backward implementation for the layernorm-MLP layer. - # set rng state for fwd - torch.set_rng_state(ctx.cpu_rng_state) - _set_cuda_rng_state(ctx.cuda_rng_state) + Caller must have populated ``args.grad_output`` / ``args.grad_ln_out`` and + the saved-tensor fields before invocation. Returns ``(dgrad, dgamma, + dbeta, fc1_wgrad, fc1_bias_grad, fc2_wgrad, fc2_bias_grad)``. + """ + with get_nvtx_range_context("_LayerNormMLP_backward"): + inputmat = args.inputmat + ln_weight = args.ln_weight + ln_out = args.ln_out + fc1_weight = args.fc1_weight_fp8 + origin_fc1_weight = args.fc1_weight + fc1_bias = args.fc1_bias + fc1_out = args.fc1_out + fc1_out_without_bias = args.fc1_out_without_bias + act_out = args.act_out + fc2_weight = args.fc2_weight_fp8 + origin_fc2_weight = args.fc2_weight + fc2_bias = args.fc2_bias + mu = args.mu + rsigma = args.rsigma + grad_output_arg = args.grad_output + + # Restore origin weights from weakrefs + # Only needed when fuse_wgrad_accumulation is enabled. + fc1_weight_python_object = None + fc2_weight_python_object = None + fc1_weight_main_grad = None + fc2_weight_main_grad = None + if args.fuse_wgrad_accumulation: + fc1_weight_ref = args.fc1_weight_ref + fc2_weight_ref = args.fc2_weight_ref + args.fc1_weight_ref = None + args.fc2_weight_ref = None + fc1_weight_python_object = fc1_weight_ref() if fc1_weight_ref is not None else None + fc2_weight_python_object = fc2_weight_ref() if fc2_weight_ref is not None else None + if args.fc1_weight_requires_grad: + assert ( + fc1_weight_python_object is not None + ), "fc1_weight was removed while fuse_wgrad_accumulation=True" + fc1_weight_main_grad = args.fc1_main_grad_func() + fc1_weight_python_object.main_grad = fc1_weight_main_grad + if args.fc2_weight_requires_grad: + assert ( + fc2_weight_python_object is not None + ), "fc2_weight was removed while fuse_wgrad_accumulation=True" + fc2_weight_main_grad = args.fc2_main_grad_func() + fc2_weight_python_object.main_grad = fc2_weight_main_grad + + # TODO: Fix this # pylint: disable=fixme + # Gather saved autograd context tensors when running with FSDP + # NOTE: weight_fp8 = weight when fp8 == False and torch.disttributed.FSDP already + # shards/unshards the base weights so we don't do it ourselves + # _fsdp_gather_tensors( + # args.fsdp_group, + # args.fsdp_shapes, + # mu, + # rsigma, + # ln_out, + # fc1_out_without_bias if bias_gelu_nvfusion else fc1_out,, + # gelu_out, + # fc1_weight_fp8 if args.fp8 and not isinstance(fc1_weight, Float8Tensor) else None, + # fc2_weight_fp8 if args.fp8 and not isinstance(fc2_weight, Float8Tensor) else None, + # ) - # Unpack saved tensors and pass None for weight workspaces (recomputed from scratch) - ( - inp_r, - ln_weight_r, - ln_bias_r, - fc1_weight_r, - fc1_bias_r, - fc2_weight_r, - fc2_bias_r, - ) = tensors - out = _LayerNormMLP._forward( # recompute - ctx, - inp_r, - ln_weight_r, - ln_bias_r, - fc1_weight_r, - None, - fc1_bias_r, - fc2_weight_r, - None, - fc2_bias_r, - tuple(ctx.other_args.values()), + # Choose whether to use GEMM kernel with split accumulator + dgrad_use_split_accumulator = args.dgrad_use_split_accumulator + wgrad_use_split_accumulator = args.wgrad_use_split_accumulator + + # No need to do bulk DGRAD/WGRAD overlap if WGRAD is not required + ub_bulk_dgrad = args.fc1_weight_requires_grad and args.ub_bulk_dgrad + ub_bulk_wgrad = args.fc1_weight_requires_grad and args.ub_bulk_wgrad + + # Configure quantizer for FC2 grad output tensor + # Note: dgrad GEMM requires row-wise usage, wgrad GEMM + # requires column-wise usage + if args.fc2_grad_output_quantizer is not None: + quantizer = args.fc2_grad_output_quantizer + quantizer.set_usage(rowwise=True, columnwise=True) + if args.ub_overlap_ag: + # Userbuffers only supports communication for one + # tensor usage at a time. Configure quantizer with + # usage for only dgrad GEMM. + quantizer.set_usage(columnwise=False) + # Amax reduction group for FC2 grad output (row-parallel sequence parallel) + set_quantizer_amax_reduction_group( + quantizer, + args.tp_group if (args.sequence_parallel and args.set_parallel_mode) else None, ) - FP8GlobalStateManager.set_autocast_state(final_autocast_state) - if ( - ctx.other_args["fp8"] - and FP8GlobalStateManager.get_fp8_recipe().__class__.__name__ == "DelayedScaling" - ): - FP8GlobalStateManager.restore_fp8_meta_tensors( - ctx.other_args["fp8_meta"] - ) # restore quantizers - - # set rng state for fwd - torch.set_rng_state(final_cpu_rng_state) - _set_cuda_rng_state(final_cuda_rng_state) - - return out - - # load from saved (return ctx is just because the other branch does too) - return tuple([ctx] + tensors) + # Prepare FC2 grad output tensor + # Note: Cast to expected dtype and perform tensor-parallel communication + ub_obj_fc2_dgrad = None + if args.ub_overlap_ag: + ub_obj_fc2_dgrad = get_ub("fc2_dgrad", args.fp8) + args.ub_obj_gradout = ub_obj_fc2_dgrad + ( + grad_output, + fc2_bias_grad, + ) = TransformerEngineBaseModule.grad_output_preprocess( + args, grad_output_arg, True, args.fc2_grad_output_quantizer + ) - @staticmethod - def backward( - ctx, *grad_outputs: Tuple[torch.Tensor, ...] - ) -> Tuple[Union[torch.Tensor, None], ...]: - # pylint: disable=missing-function-docstring - with get_nvtx_range_context("_LayerNormMLP_backward"): - ( # pylint: disable=unbalanced-tuple-unpacking - ctx, - inputmat, - ln_weight, - ln_out, - fc1_weight, - origin_fc1_weight, - fc1_bias, - fc1_out, - fc1_out_without_bias, - act_out, - fc2_weight, - origin_fc2_weight, - fc2_bias, - mu, - rsigma, - ) = _LayerNormMLP._recompute(ctx) - - # Restore origin weights from weakrefs - # Only needed when fuse_wgrad_accumulation is enabled. - fc1_weight_python_object = None - fc2_weight_python_object = None - fc1_weight_main_grad = None - fc2_weight_main_grad = None - if ctx.fuse_wgrad_accumulation: - fc1_weight_python_object_ref = getattr(ctx, "fc1_weight_python_object_ref", None) - fc2_weight_python_object_ref = getattr(ctx, "fc2_weight_python_object_ref", None) - ctx.fc1_weight_python_object_ref = None - ctx.fc2_weight_python_object_ref = None - fc1_weight_python_object = ( - fc1_weight_python_object_ref() - if fc1_weight_python_object_ref is not None - else None - ) - fc2_weight_python_object = ( - fc2_weight_python_object_ref() - if fc2_weight_python_object_ref is not None - else None - ) - if ctx.fc1_weight_requires_grad: - assert ( - fc1_weight_python_object is not None - ), "fc1_weight was removed while fuse_wgrad_accumulation=True" - fc1_weight_main_grad = ctx.fc1_main_grad_func() - fc1_weight_python_object.main_grad = fc1_weight_main_grad - if ctx.fc2_weight_requires_grad: - assert ( - fc2_weight_python_object is not None - ), "fc2_weight was removed while fuse_wgrad_accumulation=True" - fc2_weight_main_grad = ctx.fc2_main_grad_func() - fc2_weight_python_object.main_grad = fc2_weight_main_grad - - # TODO: Fix this # pylint: disable=fixme - # Gather saved autograd context tensors when running with FSDP - # NOTE: weight_fp8 = weight when ctx.fp8 == False and torch.disttributed.FSDP already - # shards/unshards the base weights so we don't do it ourselves - # _fsdp_gather_tensors( - # ctx.fsdp_group, - # ctx.fsdp_shapes, - # mu, - # rsigma, - # ln_out, - # fc1_out_without_bias if bias_gelu_nvfusion else fc1_out,, - # gelu_out, - # fc1_weight_fp8 if ctx.fp8 and not isinstance(fc1_weight, Float8Tensor) else None, - # fc2_weight_fp8 if ctx.fp8 and not isinstance(fc2_weight, Float8Tensor) else None, - # ) - - # Choose whether to use GEMM kernel with split accumulator - dgrad_use_split_accumulator = _2X_ACC_DGRAD - wgrad_use_split_accumulator = _2X_ACC_WGRAD - if ctx.fp8: - recipe = ctx.fp8_recipe - if hasattr(recipe, "fp8_gemm_dgrad"): - dgrad_use_split_accumulator = recipe.fp8_gemm_dgrad.use_split_accumulator - if hasattr(recipe, "fp8_gemm_wgrad"): - wgrad_use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator - - # No need to do bulk DGRAD/WGRAD overlap if WGRAD is not required - ctx.ub_bulk_dgrad = ctx.fc1_weight_requires_grad and ctx.ub_bulk_dgrad - ctx.ub_bulk_wgrad = ctx.fc1_weight_requires_grad and ctx.ub_bulk_wgrad - - # Configure quantizer for FC2 grad output tensor - # Note: dgrad GEMM requires row-wise usage, wgrad GEMM - # requires column-wise usage - if ctx.fc2_grad_output_quantizer is not None: - quantizer = ctx.fc2_grad_output_quantizer - quantizer.set_usage(rowwise=True, columnwise=True) - if ctx.ub_overlap_ag: - # Userbuffers only supports communication for one - # tensor usage at a time. Configure quantizer with - # usage for only dgrad GEMM. - quantizer.set_usage(columnwise=False) - # Amax reduction group for FC2 grad output (row-parallel sequence parallel) - set_quantizer_amax_reduction_group( + # Launch tensor-parallel communication for FC1 GEMM input + ln_out_total = None + ln_out_total_work = None + ub_obj_fc1_dgrad = None + if args.fc1_weight_requires_grad and args.tensor_parallel and args.sequence_parallel: + quantizer = None + if args.fp8 or args.debug: + quantizer = args.fc1_input_quantizer + set_quantizer_usage_for_wgrad_all_gather(quantizer) + if ub_bulk_dgrad: + ub_obj_fc1_dgrad = get_ub("fc1_dgrad", args.fp8) + ln_out_total, _ = fill_userbuffers_buffer_for_all_gather( + ub_obj_fc1_dgrad, + ln_out, quantizer, - ctx.tp_group if (ctx.sequence_parallel and ctx.set_parallel_mode) else None, + args.tp_group, ) - - # Prepare FC2 grad output tensor - # Note: Cast to expected dtype and perform tensor-parallel communication - ub_obj_fc2_dgrad = None - if ctx.ub_overlap_ag: - ub_obj_fc2_dgrad = get_ub("fc2_dgrad", ctx.fp8) - ctx.ub_obj_gradout = ub_obj_fc2_dgrad - ( - grad_output, - fc2_bias_grad, - ) = TransformerEngineBaseModule.grad_output_preprocess( - ctx, grad_outputs[0], True, ctx.fc2_grad_output_quantizer - ) - - # Launch tensor-parallel communication for FC1 GEMM input - ln_out_total = None - ln_out_total_work = None - ub_obj_fc1_dgrad = None - if ctx.fc1_weight_requires_grad and ctx.tensor_parallel and ctx.sequence_parallel: - quantizer = None - if ctx.fp8 or ctx.debug: - quantizer = ctx.fc1_input_quantizer - 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( - ub_obj_fc1_dgrad, - ln_out, - quantizer, - ctx.tp_group, - ) - else: - ln_out_total, ln_out_total_work = gather_along_first_dim( - ln_out, - ctx.tp_group, - async_op=True, - quantizer=quantizer, - ) else: - ln_out_total = ln_out - - # Check whether to output wgrad GEMM directly into main grad - if ctx.is_first_microbatch is not None: - accumulate_wgrad_into_param_main_grad = ( - ctx.fuse_wgrad_accumulation and not ctx.is_first_microbatch + ln_out_total, ln_out_total_work = gather_along_first_dim( + ln_out, + args.tp_group, + async_op=True, + quantizer=quantizer, ) - else: - accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation - - # -------------------------------------------------- - # FC2 DGRAD - # -------------------------------------------------- - - # There are 6 possible fusion paths - # 1 high-precision bias_gelu_fusion: gemm, FC1_bias + gelu, - # 2 high-precision fc2_dgrad_gemm_gelu_fusion: gemm + gelu, FC1_bias + quantize - # 3 fp8 activation+bias+quantize fusion: gemm, activation + FC1_bias + quantize - # 4 fp8 bias+quantize fusion: gemm, activation, FC1_bias + quantize - # 5 high-precision unfused: gemm, activation, FC1_bias + FC1_gemm - # 6 fp8 unfused: gemm, activation, FC1_bias + FC1_gemm - fc2_dgrad_gemm_gelu_fusion = ( - not ctx.fp8 - and (ctx.activation == "gelu") - and (not ctx.bias_gelu_fusion) - and (not ctx.debug) + else: + ln_out_total = ln_out + + # Check whether to output wgrad GEMM directly into main grad + if args.is_first_microbatch is not None: + accumulate_wgrad_into_param_main_grad = ( + args.fuse_wgrad_accumulation and not args.is_first_microbatch ) + else: + accumulate_wgrad_into_param_main_grad = args.fuse_wgrad_accumulation + + # -------------------------------------------------- + # FC2 DGRAD + # -------------------------------------------------- + + # There are 6 possible fusion paths + # 1 high-precision bias_gelu_fusion: gemm, FC1_bias + gelu, + # 2 high-precision fc2_dgrad_gemm_gelu_fusion: gemm + gelu, FC1_bias + quantize + # 3 fp8 activation+bias+quantize fusion: gemm, activation + FC1_bias + quantize + # 4 fp8 bias+quantize fusion: gemm, activation, FC1_bias + quantize + # 5 high-precision unfused: gemm, activation, FC1_bias + FC1_gemm + # 6 fp8 unfused: gemm, activation, FC1_bias + FC1_gemm + fc2_dgrad_gemm_gelu_fusion = ( + not args.fp8 + and (args.activation == "gelu") + and (not args.bias_gelu_fusion) + and (not args.debug) + ) - # FSDP2: Re-create workspace from all-gathered weight when - # workspace was not saved to avoid forward memory - # accumulation. (Issue #2681) - if fc2_weight is None: - if isinstance(origin_fc2_weight, QuantizedTensorStorage): - fc2_weight = origin_fc2_weight - elif ctx.fc2_weight_quantizer is not None: - ctx.fc2_weight_quantizer.set_usage(rowwise=True, columnwise=True) - fc2_weight = ctx.fc2_weight_quantizer(origin_fc2_weight) - - # Make sure required data is available - if isinstance(grad_output, QuantizedTensorStorage): - grad_output.update_usage(rowwise_usage=True) - if ctx.fc2_weight_quantizer is not None and isinstance( - fc2_weight, QuantizedTensorStorage - ): - fc2_weight.update_usage(columnwise_usage=True) + # FSDP2: Re-create workspace from all-gathered weight when + # workspace was not saved to avoid forward memory + # accumulation. (Issue #2681) + if fc2_weight is None: + if isinstance(origin_fc2_weight, QuantizedTensorStorage): + fc2_weight = origin_fc2_weight + elif args.fc2_weight_quantizer is not None: + args.fc2_weight_quantizer.set_usage(rowwise=True, columnwise=True) + fc2_weight = args.fc2_weight_quantizer(origin_fc2_weight) + + # Make sure required data is available + if isinstance(grad_output, QuantizedTensorStorage): + grad_output.update_usage(rowwise_usage=True) + if args.fc2_weight_quantizer is not None and isinstance(fc2_weight, QuantizedTensorStorage): + fc2_weight.update_usage(columnwise_usage=True) + + # Perform GEMM + gemm_output, *_ = general_gemm( + fc2_weight, + grad_output, + layout="NN", + grad=True, + quantization_params=( + args.fc1_grad_input_quantizer if fc2_dgrad_gemm_gelu_fusion or args.debug else None + ), # high precision to activation + out_dtype=args.activation_dtype, + gelu=fc2_dgrad_gemm_gelu_fusion, + gelu_in=fc1_out if fc2_dgrad_gemm_gelu_fusion else None, + use_split_accumulator=dgrad_use_split_accumulator, + ub=ub_obj_fc2_dgrad, + ub_type=tex.CommOverlapType.AG if args.ub_overlap_ag else None, + ) - # Perform GEMM - gemm_output, *_ = general_gemm( - fc2_weight, + # FSDP2: Clear columnwise/transpose caches after FC2 dgrad GEMM + # to prevent them from persisting on the all-gathered buffer. + # Uses is_fsdp2 (not fsdp2_skip_columnwise) so cleanup runs + # even when backward follows gradient-checkpoint recomputation. + # (Issues #2681, #2717) + if args.is_fsdp2 and isinstance(fc2_weight, QuantizedTensorStorage): + clear_columnwise_cache(fc2_weight) + + # Prepare input grad tensor + dact = None + fc2_dgrad = None + if fc2_dgrad_gemm_gelu_fusion: + dact = gemm_output + else: + fc2_dgrad = gemm_output + + # -------------------------------------------------- + # Finished FC2 DGRAD... + # -------------------------------------------------- + + # cuBLASMp's AG+GEMM consumes the gathered grad_output inline and + # does not preserve it for fc2_wgrad. Userbuffers leaves the + # gathered tensor in its persistent buffer; cuBLASMp does not, so + # we gather here. Route through the same FP8-aware all-gather as + # the non-overlap path in + # ``TransformerEngineBaseModule.grad_output_preprocess`` by passing + # the grad_output quantizer. Per-tensor FP8 can reconstruct columnwise + # data from the gathered rowwise data; MXFP8 must instead quantize + # the original gradient columnwise to avoid double quantization. + if ( + args.fc2_weight_requires_grad + and args.ub_overlap_ag + and args.ub_obj_gradout is not None + and args.ub_obj_gradout.with_cublasmp() + ): + if args.fc2_grad_output_quantizer is not None: + set_quantizer_usage_for_wgrad_all_gather(args.fc2_grad_output_quantizer) + if isinstance(args.fc2_grad_output_quantizer, MXFP8Quantizer): + grad_output = grad_output_arg.reshape(-1, grad_output_arg.shape[-1]).contiguous() + grad_output, _ = gather_along_first_dim( grad_output, - layout="NN", - grad=True, - quantization_params=( - ctx.fc1_grad_input_quantizer - if fc2_dgrad_gemm_gelu_fusion or ctx.debug - else None - ), # high precision to activation - out_dtype=ctx.activation_dtype, - gelu=fc2_dgrad_gemm_gelu_fusion, - gelu_in=fc1_out if fc2_dgrad_gemm_gelu_fusion else None, - use_split_accumulator=dgrad_use_split_accumulator, - ub=ub_obj_fc2_dgrad, - ub_type=tex.CommOverlapType.AG if ctx.ub_overlap_ag else None, + args.tp_group, + quantizer=args.fc2_grad_output_quantizer, ) - # FSDP2: Clear columnwise/transpose caches after FC2 dgrad GEMM - # to prevent them from persisting on the all-gathered buffer. - # Uses is_fsdp2 (not fsdp2_skip_columnwise) so cleanup runs - # even when backward follows gradient-checkpoint recomputation. - # (Issues #2681, #2717) - if getattr(ctx, "is_fsdp2", False) and isinstance(fc2_weight, QuantizedTensorStorage): - clear_columnwise_cache(fc2_weight) - - # Prepare input grad tensor - dact = None - fc2_dgrad = None - if fc2_dgrad_gemm_gelu_fusion: - dact = gemm_output - else: - fc2_dgrad = gemm_output - - # -------------------------------------------------- - # Finished FC2 DGRAD... - # -------------------------------------------------- - - # cuBLASMp's AG+GEMM consumes the gathered grad_output inline and - # does not preserve it for fc2_wgrad. Userbuffers leaves the - # gathered tensor in its persistent buffer; cuBLASMp does not, so - # we gather here. Route through the same FP8-aware all-gather as - # the non-overlap path in - # ``TransformerEngineBaseModule.grad_output_preprocess`` by passing - # the grad_output quantizer. Columnwise data needed for fc2_wgrad - # is produced by ``update_usage(columnwise_usage=True)`` further - # below. + # -------------------------------------------------- + # FC2 WGRAD + # -------------------------------------------------- + + fc2_wgrad = None + if args.fc2_weight_requires_grad: + # Prepare grad output tensor + # Note: Synchronize tensor-parallel communication and + # make sure required data is available if ( - ctx.fc2_weight_requires_grad - and ctx.ub_overlap_ag - and ctx.ub_obj_gradout is not None - and ctx.ub_obj_gradout.with_cublasmp() + args.ub_overlap_ag + and isinstance(args.fc2_grad_output_quantizer, MXFP8Quantizer) + and not ub_obj_fc2_dgrad.with_cublasmp() ): - if ctx.fc2_grad_output_quantizer is not None: - set_quantizer_usage_for_wgrad_all_gather(ctx.fc2_grad_output_quantizer) - grad_output, _ = gather_along_first_dim( - grad_output, - ctx.tp_group, - quantizer=ctx.fc2_grad_output_quantizer, + # UB does not support pipelined overlapping grad output + # all-gather with wgrad GEMM. Also, we can't + # convert row-scaled MXFP8 to column-scaled, so we + # can't reuse the grad output that was gathered + # for the dgrad GEMM. We work around by explicitly + # overlapping the AG operation with the dgrad GEMM. + + # Get the communication stream from the dgrad GEMM to use for the AG + dgrad_send_stream, dgrad_recv_stream = ub_obj_fc2_dgrad.get_communication_stream() + + ub_obj_fc2_wgrad = get_ub("fc2_wgrad", args.fp8) + + args.fc2_grad_output_quantizer.set_usage(rowwise=False, columnwise=True) + + # We use the send stream to copy into the userbuffers. + # This is the same stream that we will use to access the data in the AG, + # so we dont need to add any syncs yet. + with torch.cuda.stream(dgrad_send_stream): + grad_output, _ = fill_userbuffers_buffer_for_all_gather( + ub_obj_fc2_wgrad, + grad_output_arg, + args.fc2_grad_output_quantizer, + args.tp_group, + ) + + # Allgather grad_outputs[0] using the dgrad streams so we can overlap with the fc2_dgrad gemm + tex.bulk_overlap_ag_with_external_gemm( + ub_obj_fc2_wgrad, dgrad_send_stream, dgrad_recv_stream ) - # -------------------------------------------------- - # FC2 WGRAD - # -------------------------------------------------- + # Prepare input tensor + # Note: Synchronize tensor-parallel communication and + # make sure required data is available + if args.fp8 or args.debug: + if isinstance(act_out, QuantizedTensorStorage): + act_out.update_usage(columnwise_usage=True) + else: + args.fc2_input_quantizer.set_usage(rowwise=False, columnwise=True) + act_out = args.fc2_input_quantizer(act_out) - fc2_wgrad = None - if ctx.fc2_weight_requires_grad: - # Prepare grad output tensor - # Note: Synchronize tensor-parallel communication and - # make sure required data is available - if ctx.ub_overlap_ag and isinstance(ctx.fc2_grad_output_quantizer, MXFP8Quantizer): - # UB does not support pipelined overlapping grad output - # all-gather with wgrad GEMM. Also, we can't - # convert row-scaled MXFP8 to column-scaled, so we - # can't reuse the grad output that was gathered - # for the dgrad GEMM. We work around by explicitly - # overlapping the AG operation with the dgrad GEMM. - - # Get the communication stream from the dgrad GEMM to use for the AG - dgrad_send_stream, dgrad_recv_stream = ( - ub_obj_fc2_dgrad.get_communication_stream() - ) + if args.fp8 or args.debug: + if isinstance(grad_output, QuantizedTensorStorage): + grad_output.update_usage(columnwise_usage=True) + else: + args.fc2_grad_output_quantizer.set_usage(rowwise=False, columnwise=True) + grad_output = args.fc2_grad_output_quantizer(grad_output) + + # Whether to set grad arg in general_gemm + grad_arg = True + if args.fp8 and args.fp8_recipe.float8_block_scaling(): + grad_arg = False + + # Arguments to include in wgrad GEMM closure + fc2_wgrad_gemm_kwargs = { + "out_dtype": ( + fc2_weight_main_grad.dtype + if args.fuse_wgrad_accumulation + else args.activation_dtype + ), + "quantization_params": args.fc2_grad_weight_quantizer, # wgrad in high precision + "accumulate": ( + accumulate_wgrad_into_param_main_grad + if not args.fc2_weight_overwrites_main_grad + else False + ), + "layout": "NT", + "out": fc2_weight_main_grad if args.fuse_wgrad_accumulation else None, + "bias": fc2_bias if fc2_bias is not None and fc2_bias_grad is None else None, + "use_split_accumulator": wgrad_use_split_accumulator, + "grad": grad_arg, + } - ub_obj_fc2_wgrad = get_ub("fc2_wgrad", ctx.fp8) + def fc2_wgrad_gemm( + x: torch.Tensor, + dy: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Perform FC2 WGRAD GEMM - ctx.fc2_grad_output_quantizer.set_usage(rowwise=False, columnwise=True) + May be called outside of this function to enable + some advanced communication/compute overlapping. - # We use the send stream to copy into the userbuffers. - # This is the same stream that we will use to access the data in the AG, - # so we dont need to add any syncs yet. - with torch.cuda.stream(dgrad_send_stream): - grad_output, _ = fill_userbuffers_buffer_for_all_gather( - ub_obj_fc2_wgrad, - grad_outputs[0], - ctx.fc2_grad_output_quantizer, - ctx.tp_group, - ) + """ + dw, db, *_ = general_gemm(x, dy, **fc2_wgrad_gemm_kwargs) + return dw, db - # Allgather grad_outputs[0] using the dgrad streams so we can overlap with the fc2_dgrad gemm - tex.bulk_overlap_ag_with_external_gemm( - ub_obj_fc2_wgrad, dgrad_send_stream, dgrad_recv_stream - ) + # Choose whether to call wgrad GEMM now or delay + if args.wgrad_store is not None and args.wgrad_store.delay_wgrad_compute(): + args.wgrad_store.put([act_out, grad_output], fc2_wgrad_gemm) + else: - # Prepare input tensor - # Note: Synchronize tensor-parallel communication and - # make sure required data is available - if ctx.fp8 or ctx.debug: - if isinstance(act_out, QuantizedTensorStorage): - act_out.update_usage(columnwise_usage=True) - else: - ctx.fc2_input_quantizer.set_usage(rowwise=False, columnwise=True) - act_out = ctx.fc2_input_quantizer(act_out) - - if ctx.fp8 or ctx.debug: - if isinstance(grad_output, QuantizedTensorStorage): - grad_output.update_usage(columnwise_usage=True) - else: - ctx.fc2_grad_output_quantizer.set_usage(rowwise=False, columnwise=True) - grad_output = ctx.fc2_grad_output_quantizer(grad_output) - - # Whether to set grad arg in general_gemm - grad_arg = True - if ctx.fp8 and ctx.fp8_recipe.float8_block_scaling(): - grad_arg = False - - # Arguments to include in wgrad GEMM closure - fc2_wgrad_gemm_kwargs = { - "out_dtype": ( - fc2_weight_main_grad.dtype - if ctx.fuse_wgrad_accumulation - else ctx.activation_dtype - ), - "quantization_params": ctx.fc2_grad_weight_quantizer, # wgrad in high precision - "accumulate": ( - accumulate_wgrad_into_param_main_grad - if not getattr(ctx, "fc2_weight_overwrites_main_grad", False) - else False - ), - "layout": "NT", - "out": fc2_weight_main_grad if ctx.fuse_wgrad_accumulation else None, - "bias": fc2_bias if fc2_bias is not None and fc2_bias_grad is None else None, - "use_split_accumulator": wgrad_use_split_accumulator, - "grad": grad_arg, - } - - def fc2_wgrad_gemm( - x: torch.Tensor, - dy: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Perform FC2 WGRAD GEMM - - May be called outside of this function to enable - some advanced communication/compute overlapping. - - """ - dw, db, *_ = general_gemm(x, dy, **fc2_wgrad_gemm_kwargs) - return dw, db - - # Choose whether to call wgrad GEMM now or delay - if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute(): - ctx.wgrad_store.put([act_out, grad_output], fc2_wgrad_gemm) - else: + # Call wgrad GEMM now + fc2_wgrad, fc2_bias_grad_ = fc2_wgrad_gemm(act_out, grad_output) - # Call wgrad GEMM now - fc2_wgrad, fc2_bias_grad_ = fc2_wgrad_gemm(act_out, grad_output) - - # Update grad bias if needed - if fc2_bias_grad is None: - if ( - ctx.fp8 - and ctx.fp8_recipe.float8_block_scaling() - and fc2_bias is not None - ): - # BGRAD not fused with GEMM for float8 blockwise gemm. - fc2_bias_grad_ = act_out.view(-1, act_out.shape[-1]).sum(dim=0) - fc2_bias_grad = fc2_bias_grad_ - del fc2_bias_grad_ - - # Deallocate input tensor if permitted - if ctx.wgrad_store is not None and not ctx.wgrad_store.delay_wgrad_compute(): - clear_tensor_data(act_out) - - # -------------------------------------------------- - # Finished FC2 WGRAD... - # -------------------------------------------------- - - # bias computation - act_params = ctx.activation_params or {} - fc1_bias_grad = None - fuse_gemm_and_bias_fc1_wgrad = False - if ctx.fc1_grad_output_quantizer is not None: - ctx.fc1_grad_output_quantizer.set_usage(rowwise=True, columnwise=True) - if ctx.bias_gelu_fusion: - # Fusion: gemm, bias + gelu - assert ctx.activation == "gelu" - assert not ctx.fp8 - fc1_bias_grad, dact = bgrad_dgelu_fused(fc2_dgrad, fc1_out_without_bias, fc1_bias) - if ctx.fc1_grad_output_quantizer is not None: - dact = ctx.fc1_grad_output_quantizer(dact) - elif ctx.debug: - dact_func = _act_func(ctx.activation)[1] - dact = dact_func(fc2_dgrad, fc1_out.to(ctx.activation_dtype), None, **act_params) - fc1_bias_grad = dact.sum(dim=0) - dact = ctx.fc1_grad_output_quantizer(dact) - elif ( - _act_func(ctx.activation, ctx.fp8_recipe if ctx.fp8 else None)[2] is not None - and ctx.fp8 - ): - # Fusion: gemm, bias + gelu + quantize - dbias_dact_quantize_func = _act_func( - ctx.activation, ctx.fp8_recipe if ctx.fp8 else None - )[2] - fc1_bias_grad, dact = dbias_dact_quantize_func( - fc2_dgrad, - fc1_out.to(ctx.activation_dtype), - ctx.fc1_grad_output_quantizer, - **act_params, - ) # quantize bgrad gelu fused - else: - # Fusion: gemm + gelu, - if not fc2_dgrad_gemm_gelu_fusion: - activation_func_bwd = _act_func( - ctx.activation, ctx.fp8_recipe if ctx.fp8 else None - )[1] - dact = activation_func_bwd( - fc2_dgrad, fc1_out.to(ctx.activation_dtype), None, **act_params - ) # activation in high precision - - 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, IdentityQuantizer), - ) - or ctx.fp8_recipe.custom() - ): - fc1_bias_grad = dact.view(-1, dact.shape[-1]).sum(dim=0) - dact = ctx.fc1_grad_output_quantizer(dact) - else: - fc1_bias_grad, dact = tex.bgrad_quantize( - dact, ctx.fc1_grad_output_quantizer - ) - else: - fuse_gemm_and_bias_fc1_wgrad = ( - True # fc1_bias_grad is computed later, fused with wgrad gemm for the FC1 + # Update grad bias if needed + if fc2_bias_grad is None: + if args.fp8 and args.fp8_recipe.float8_block_scaling() and fc2_bias is not None: + # BGRAD not fused with GEMM for float8 blockwise gemm. + fc2_bias_grad_ = act_out.view(-1, act_out.shape[-1]).sum(dim=0) + fc2_bias_grad = fc2_bias_grad_ + del fc2_bias_grad_ + + # Deallocate input tensor if permitted + if args.wgrad_store is not None and not args.wgrad_store.delay_wgrad_compute(): + clear_tensor_data(act_out) + + # -------------------------------------------------- + # Finished FC2 WGRAD... + # -------------------------------------------------- + + # bias computation + act_params = args.activation_params or {} + fc1_bias_grad = None + fuse_gemm_and_bias_fc1_wgrad = False + if args.fc1_grad_output_quantizer is not None: + args.fc1_grad_output_quantizer.set_usage(rowwise=True, columnwise=True) + if args.bias_gelu_fusion: + # Fusion: gemm, bias + gelu + assert args.activation == "gelu" + assert not args.fp8 + fc1_bias_grad, dact = bgrad_dgelu_fused(fc2_dgrad, fc1_out_without_bias, fc1_bias) + if args.fc1_grad_output_quantizer is not None: + dact = args.fc1_grad_output_quantizer(dact) + elif args.debug: + dact_func = _act_func(args.activation)[1] + dact = dact_func(fc2_dgrad, fc1_out.to(args.activation_dtype), None, **act_params) + fc1_bias_grad = dact.sum(dim=0) + dact = args.fc1_grad_output_quantizer(dact) + elif ( + _act_func(args.activation, args.fp8_recipe if args.fp8 else None)[2] is not None + and args.fp8 + ): + # Fusion: gemm, bias + gelu + quantize + dbias_dact_quantize_func = _act_func( + args.activation, args.fp8_recipe if args.fp8 else None + )[2] + fc1_bias_grad, dact = dbias_dact_quantize_func( + fc2_dgrad, + fc1_out.to(args.activation_dtype), + args.fc1_grad_output_quantizer, + **act_params, + ) # quantize bgrad gelu fused + else: + # Fusion: gemm + gelu, + if not fc2_dgrad_gemm_gelu_fusion: + activation_func_bwd = _act_func( + args.activation, args.fp8_recipe if args.fp8 else None + )[1] + dact = activation_func_bwd( + fc2_dgrad, fc1_out.to(args.activation_dtype), None, **act_params + ) # activation in high precision + + if args.fp8: + # TODO float8 blockwise current scaling (as well as custom quantizers) has no bgrad fusion for now + if ( + isinstance( + args.fc1_grad_output_quantizer, + (Float8BlockQuantizer, IdentityQuantizer), ) - # it may not be calculated in case wgrad is not required. - if fc1_bias is not None: - if not ctx.fc1_weight_requires_grad and fc1_bias.requires_grad: - fc1_bias_grad = dact.sum(dim=0) - - # Overwrite data. Deleting the tensor does not release underlying memory. - clear_tensor_data(fc1_out, fc1_out_without_bias) - - # Set UB algo and UB obj for fc1_dgrad/wgrad bulk/pipelined overlap - ub_obj_fc1_dgrad = None - ub_obj_fc1_wgrad = None - ub_type_fc1_dgrad = None - ub_type_fc1_wgrad = None - fc1_dgrad_shape = [reduce(multiply_op, inputmat.shape[:-1]), inputmat.shape[-1]] - if ctx.ub_overlap_rs_dgrad: - # Overlap DGRAD+RS - ub_obj_fc1_dgrad = get_ub("fc1_dgrad", ctx.fp8) - ub_type_fc1_dgrad = tex.CommOverlapType.RS + or args.fp8_recipe.custom() + ): + fc1_bias_grad = dact.view(-1, dact.shape[-1]).sum(dim=0) + dact = args.fc1_grad_output_quantizer(dact) + else: + fc1_bias_grad, dact = tex.bgrad_quantize(dact, args.fc1_grad_output_quantizer) else: - if ctx.ub_bulk_dgrad: - # Overlap ln_out all-gather with DGRAD compute - ub_obj_fc1_dgrad = get_ub("fc1_dgrad", ctx.fp8) - ub_type_fc1_dgrad = tex.CommOverlapType.AG - if ctx.ub_bulk_wgrad: - # Overlap FC1 DGRAD reduce-scatter with WGRAD compute - ub_obj_fc1_wgrad = get_ub("fc1_wgrad", ctx.fp8) - ub_type_fc1_wgrad = tex.CommOverlapType.RS - - # -------------------------------------------------- - # FC1 DGRAD - # -------------------------------------------------- - - # FSDP2: Re-create workspace from all-gathered weight when - # workspace was not saved. (Issue #2681) - if fc1_weight is None: - if isinstance(origin_fc1_weight, QuantizedTensorStorage): - fc1_weight = origin_fc1_weight - elif ctx.fc1_weight_quantizer is not None: - ctx.fc1_weight_quantizer.set_usage(rowwise=True, columnwise=True) - fc1_weight = ctx.fc1_weight_quantizer(origin_fc1_weight) - - # Make sure required data is available - if ctx.fc1_weight_quantizer is not None and isinstance( - fc1_weight, QuantizedTensorStorage - ): - fc1_weight.update_usage(columnwise_usage=True) + fuse_gemm_and_bias_fc1_wgrad = ( + True # fc1_bias_grad is computed later, fused with wgrad gemm for the FC1 + ) + # it may not be calculated in case wgrad is not required. + if fc1_bias is not None: + if not args.fc1_weight_requires_grad and args.fc1_bias_requires_grad: + fc1_bias_grad = dact.sum(dim=0) + + # Overwrite data. Deleting the tensor does not release underlying memory. + clear_tensor_data(fc1_out, fc1_out_without_bias) + + # Set UB algo and UB obj for fc1_dgrad/wgrad bulk/pipelined overlap + ub_obj_fc1_dgrad = None + ub_obj_fc1_wgrad = None + ub_type_fc1_dgrad = None + ub_type_fc1_wgrad = None + fc1_dgrad_shape = [reduce(multiply_op, inputmat.shape[:-1]), inputmat.shape[-1]] + if args.ub_overlap_rs_dgrad: + # Overlap DGRAD+RS + ub_obj_fc1_dgrad = get_ub("fc1_dgrad", args.fp8) + ub_type_fc1_dgrad = tex.CommOverlapType.RS + else: + if ub_bulk_dgrad: + # Overlap ln_out all-gather with DGRAD compute + ub_obj_fc1_dgrad = get_ub("fc1_dgrad", args.fp8) + ub_type_fc1_dgrad = tex.CommOverlapType.AG + if ub_bulk_wgrad: + # Overlap FC1 DGRAD reduce-scatter with WGRAD compute + ub_obj_fc1_wgrad = get_ub("fc1_wgrad", args.fp8) + ub_type_fc1_wgrad = tex.CommOverlapType.RS + + # -------------------------------------------------- + # FC1 DGRAD + # -------------------------------------------------- + + # FSDP2: Re-create workspace from all-gathered weight when + # workspace was not saved. (Issue #2681) + if fc1_weight is None: + if isinstance(origin_fc1_weight, QuantizedTensorStorage): + fc1_weight = origin_fc1_weight + elif args.fc1_weight_quantizer is not None: + args.fc1_weight_quantizer.set_usage(rowwise=True, columnwise=True) + fc1_weight = args.fc1_weight_quantizer(origin_fc1_weight) + + # Make sure required data is available + if args.fc1_weight_quantizer is not None and isinstance(fc1_weight, QuantizedTensorStorage): + fc1_weight.update_usage(columnwise_usage=True) + + # Output buffers for Userbuffers reduce-scatter + gemm_out = None + reduce_scatter_out = None + if args.ub_overlap_rs_dgrad: + reduce_scatter_out = torch.empty( + fc1_dgrad_shape, dtype=args.activation_dtype, device="cuda" + ) + if ub_bulk_wgrad: + gemm_out = ub_obj_fc1_wgrad.get_buffer(local_chunk=False) + + # dgrad GEMM + gemm_out, *_, reduce_scatter_out = general_gemm( + fc1_weight, + dact, + out=gemm_out, + out_dtype=args.activation_dtype, + quantization_params=args.fc1_grad_input_quantizer, + layout="NN", + grad=True, + use_split_accumulator=dgrad_use_split_accumulator, + ub=ub_obj_fc1_dgrad, + ub_type=ub_type_fc1_dgrad, + extra_output=reduce_scatter_out, + bulk_overlap=ub_bulk_dgrad, + ) + + # FSDP2: Clear columnwise/transpose caches after FC1 dgrad GEMM + # to prevent them from persisting on the all-gathered buffer. + # Uses is_fsdp2 (not fsdp2_skip_columnwise) so cleanup runs + # even when backward follows gradient-checkpoint recomputation. + # (Issues #2681, #2717) + if args.is_fsdp2 and isinstance(fc1_weight, QuantizedTensorStorage): + clear_columnwise_cache(fc1_weight) + + # Prepare grad input tensor + # Note: Perform tensor-parallel communication + fc1_dgrad = None + fc1_dgrad_work = None + if args.ub_overlap_rs_dgrad: + # cuBLASMp writes the reduce-scattered dgrad directly into the + # GEMM output tensor; Userbuffers uses the extra-output buffer. + fc1_dgrad = ( + gemm_out + if ub_obj_fc1_dgrad is not None and ub_obj_fc1_dgrad.with_cublasmp() + else reduce_scatter_out + ) + elif ub_bulk_wgrad: + fc1_dgrad = ub_obj_fc1_wgrad.get_buffer(local_chunk=True) + elif args.set_parallel_mode and not ub_bulk_wgrad: + fc1_dgrad = gemm_out + if args.sequence_parallel: + if args.return_layernorm_output and args.return_layernorm_output_gathered: + fc1_dgrad = fc1_dgrad + args.grad_ln_out.view_as(fc1_dgrad) + fc1_dgrad, fc1_dgrad_work = reduce_scatter_along_first_dim( + fc1_dgrad, + args.tp_group, + async_op=True, + ) + elif args.tensor_parallel: + fc1_dgrad, fc1_dgrad_work = allreduce(fc1_dgrad, args.tp_group, async_op=True) + else: + fc1_dgrad = gemm_out + + # -------------------------------------------------- + # Finished FC1 DGRAD... + # -------------------------------------------------- + + # -------------------------------------------------- + # FC1 WGRAD + # -------------------------------------------------- + fc1_wgrad = None + if args.fc1_weight_requires_grad: - # Output buffers for Userbuffers reduce-scatter - gemm_out = None + # Prepare input tensor + # Note: Synchronize tensor-parallel communication and + # make sure required data is available + if ln_out_total_work is not None: + ln_out_total_work.wait() + ln_out_total_work = None + if args.fp8 or args.debug: + if isinstance(ln_out_total, QuantizedTensorStorage): + ln_out_total.update_usage(columnwise_usage=True) + else: + args.fc1_input_quantizer.set_usage(rowwise=False, columnwise=True) + ln_out_total = args.fc1_input_quantizer(ln_out_total) + + # Prepare grad output tensor + # Note: Synchronize tensor-parallel communication and + # make sure required data is available + if args.fp8 or args.debug: + if isinstance(dact, QuantizedTensorStorage): + dact.update_usage(columnwise_usage=True) + else: + args.fc1_grad_output_quantizer.set_usage(rowwise=False, columnwise=True) + dact = args.fc1_grad_output_quantizer(dact) + + # Output buffer for overlapping grad input + # reduce-scatter with wgrad GEMM reduce_scatter_out = None - if ctx.ub_overlap_rs_dgrad: + if ub_bulk_wgrad and ub_obj_fc1_wgrad.is_fp8_ubuf(): reduce_scatter_out = torch.empty( - fc1_dgrad_shape, dtype=ctx.activation_dtype, device="cuda" + fc1_dgrad_shape, dtype=args.activation_dtype, device="cuda" ) - if ctx.ub_bulk_wgrad: - gemm_out = ub_obj_fc1_wgrad.get_buffer(local_chunk=False) - # dgrad GEMM - gemm_out, *_, reduce_scatter_out = general_gemm( - fc1_weight, - dact, - out=gemm_out, - out_dtype=ctx.activation_dtype, - quantization_params=ctx.fc1_grad_input_quantizer, - layout="NN", - grad=True, - use_split_accumulator=dgrad_use_split_accumulator, - ub=ub_obj_fc1_dgrad, - ub_type=ub_type_fc1_dgrad, - extra_output=reduce_scatter_out, - bulk_overlap=ctx.ub_bulk_dgrad, - ) + # Arguments to include in wgrad GEMM closure + fc1_wgrad_gemm_kwargs = { + "out_dtype": ( + fc1_weight_main_grad.dtype + if args.fuse_wgrad_accumulation + else args.activation_dtype + ), + "quantization_params": args.fc1_grad_weight_quantizer, + "accumulate": ( + accumulate_wgrad_into_param_main_grad + if not args.fc1_weight_overwrites_main_grad + else False + ), + "layout": "NT", + "out": fc1_weight_main_grad if args.fuse_wgrad_accumulation else None, + "bias": fc1_bias if fuse_gemm_and_bias_fc1_wgrad else None, + "use_split_accumulator": wgrad_use_split_accumulator, + "grad": fuse_gemm_and_bias_fc1_wgrad, + "ub": ub_obj_fc1_wgrad, + "ub_type": ub_type_fc1_wgrad, + "extra_output": reduce_scatter_out, + "bulk_overlap": ub_bulk_wgrad, + } - # FSDP2: Clear columnwise/transpose caches after FC1 dgrad GEMM - # to prevent them from persisting on the all-gathered buffer. - # Uses is_fsdp2 (not fsdp2_skip_columnwise) so cleanup runs - # even when backward follows gradient-checkpoint recomputation. - # (Issues #2681, #2717) - if getattr(ctx, "is_fsdp2", False) and isinstance(fc1_weight, QuantizedTensorStorage): - clear_columnwise_cache(fc1_weight) - - # Prepare grad input tensor - # Note: Perform tensor-parallel communication - fc1_dgrad = None - fc1_dgrad_work = None - if ctx.ub_overlap_rs_dgrad: - # cuBLASMp writes the reduce-scattered dgrad directly into the - # GEMM output tensor; Userbuffers uses the extra-output buffer. - fc1_dgrad = ( - gemm_out - if ub_obj_fc1_dgrad is not None and ub_obj_fc1_dgrad.with_cublasmp() - else reduce_scatter_out - ) - elif ctx.ub_bulk_wgrad: - fc1_dgrad = ub_obj_fc1_wgrad.get_buffer(local_chunk=True) - elif ctx.set_parallel_mode and not ctx.ub_bulk_wgrad: - fc1_dgrad = gemm_out - if ctx.sequence_parallel: - if ctx.return_layernorm_output and ctx.return_layernorm_output_gathered: - fc1_dgrad = fc1_dgrad + grad_outputs[1].view_as(fc1_dgrad) - fc1_dgrad, fc1_dgrad_work = reduce_scatter_along_first_dim( - fc1_dgrad, - ctx.tp_group, - async_op=True, + def fc1_wgrad_gemm( + x: torch.Tensor, + dy: torch.Tensor, + _is_delayed: bool = True, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Perform FC1 WGRAD GEMM + + May be called outside of this function to enable + some advanced communication/compute overlapping. + + """ + dw, db, *_ = general_gemm(x, dy, **fc1_wgrad_gemm_kwargs) + return dw, db + + # Choose whether to call wgrad GEMM now or delay + if args.wgrad_store is not None and args.wgrad_store.delay_wgrad_compute(): + if ( + fc1_wgrad_gemm_kwargs["ub"] is not None + or fc1_wgrad_gemm_kwargs["ub_type"] is not None + or fc1_wgrad_gemm_kwargs["extra_output"] is not None + or fc1_wgrad_gemm_kwargs["bulk_overlap"] + ): + raise NotImplementedError( + "Delayed weight grad computation is not supported " + "with Userbuffers (tensor-parallel communication overlapping)" ) - elif ctx.tensor_parallel: - fc1_dgrad, fc1_dgrad_work = allreduce(fc1_dgrad, ctx.tp_group, async_op=True) + args.wgrad_store.put([ln_out_total, dact], fc1_wgrad_gemm) + if fuse_gemm_and_bias_fc1_wgrad: + fc1_bias_grad = None else: - fc1_dgrad = gemm_out - # -------------------------------------------------- - # Finished FC1 DGRAD... - # -------------------------------------------------- + # Call wgrad GEMM now + fc1_wgrad_outputs = fc1_wgrad_gemm(ln_out_total, dact) + if fuse_gemm_and_bias_fc1_wgrad: + fc1_wgrad, fc1_bias_grad = fc1_wgrad_outputs + else: + fc1_wgrad, _ = fc1_wgrad_outputs - # -------------------------------------------------- - # FC1 WGRAD - # -------------------------------------------------- - fc1_wgrad = None - if ctx.fc1_weight_requires_grad: - - # Prepare input tensor - # Note: Synchronize tensor-parallel communication and - # make sure required data is available - if ln_out_total_work is not None: - ln_out_total_work.wait() - ln_out_total_work = None - if ctx.fp8 or ctx.debug: - if isinstance(ln_out_total, QuantizedTensorStorage): - ln_out_total.update_usage(columnwise_usage=True) - else: - ctx.fc1_input_quantizer.set_usage(rowwise=False, columnwise=True) - ln_out_total = ctx.fc1_input_quantizer(ln_out_total) - - # Prepare grad output tensor - # Note: Synchronize tensor-parallel communication and - # make sure required data is available - if ctx.fp8 or ctx.debug: - if isinstance(dact, QuantizedTensorStorage): - dact.update_usage(columnwise_usage=True) - else: - ctx.fc1_grad_output_quantizer.set_usage(rowwise=False, columnwise=True) - dact = ctx.fc1_grad_output_quantizer(dact) - - # Output buffer for overlapping grad input - # reduce-scatter with wgrad GEMM - reduce_scatter_out = None - if ctx.ub_bulk_wgrad and ub_obj_fc1_wgrad.is_fp8_ubuf(): - reduce_scatter_out = torch.empty( - fc1_dgrad_shape, dtype=ctx.activation_dtype, device="cuda" - ) + # Deallocate tensors if permitted + clear_tensor_data(dact) + if not args.return_layernorm_output_gathered: + clear_tensor_data(ln_out_total) - # Arguments to include in wgrad GEMM closure - fc1_wgrad_gemm_kwargs = { - "out_dtype": ( - fc1_weight_main_grad.dtype - if ctx.fuse_wgrad_accumulation - else ctx.activation_dtype - ), - "quantization_params": ctx.fc1_grad_weight_quantizer, - "accumulate": ( - accumulate_wgrad_into_param_main_grad - if not getattr(ctx, "fc1_weight_overwrites_main_grad", False) - else False - ), - "layout": "NT", - "out": fc1_weight_main_grad if ctx.fuse_wgrad_accumulation else None, - "bias": fc1_bias if fuse_gemm_and_bias_fc1_wgrad else None, - "use_split_accumulator": wgrad_use_split_accumulator, - "grad": fuse_gemm_and_bias_fc1_wgrad, - "ub": ub_obj_fc1_wgrad, - "ub_type": ub_type_fc1_wgrad, - "extra_output": reduce_scatter_out, - "bulk_overlap": ctx.ub_bulk_wgrad, - } - - def fc1_wgrad_gemm( - x: torch.Tensor, - dy: torch.Tensor, - _is_delayed: bool = True, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Perform FC1 WGRAD GEMM - - May be called outside of this function to enable - some advanced communication/compute overlapping. - - """ - dw, db, *_ = general_gemm(x, dy, **fc1_wgrad_gemm_kwargs) - return dw, db - - # Choose whether to call wgrad GEMM now or delay - if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute(): - if ( - fc1_wgrad_gemm_kwargs["ub"] is not None - or fc1_wgrad_gemm_kwargs["ub_type"] is not None - or fc1_wgrad_gemm_kwargs["extra_output"] is not None - or fc1_wgrad_gemm_kwargs["bulk_overlap"] - ): - raise NotImplementedError( - "Delayed weight grad computation is not supported " - "with Userbuffers (tensor-parallel communication overlapping)" - ) - ctx.wgrad_store.put([ln_out_total, dact], fc1_wgrad_gemm) - if fuse_gemm_and_bias_fc1_wgrad: - fc1_bias_grad = None + # Update grad input if overlapping reduce-scatter with wgrad GEMM + if ub_bulk_wgrad: + if ub_obj_fc1_wgrad.is_fp8_ubuf(): + fc1_dgrad = reduce_scatter_out else: + fc1_dgrad = ub_obj_fc1_wgrad.get_buffer(local_chunk=True).clone() - # Call wgrad GEMM now - fc1_wgrad_outputs = fc1_wgrad_gemm(ln_out_total, dact) - if fuse_gemm_and_bias_fc1_wgrad: - fc1_wgrad, fc1_bias_grad = fc1_wgrad_outputs - else: - fc1_wgrad, _ = fc1_wgrad_outputs - - # Deallocate tensors if permitted - clear_tensor_data(dact) - if not ctx.return_layernorm_output_gathered: - clear_tensor_data(ln_out_total) - - # Update grad input if overlapping reduce-scatter with wgrad GEMM - if ctx.ub_bulk_wgrad: - if ub_obj_fc1_wgrad.is_fp8_ubuf(): - fc1_dgrad = reduce_scatter_out - else: - fc1_dgrad = ub_obj_fc1_wgrad.get_buffer(local_chunk=True).clone() - - # -------------------------------------------------- - # Finished FC1 WGRAD... - # -------------------------------------------------- - - # Make sure all tensor-parallel communication is finished - if ln_out_total_work is not None: - ln_out_total_work.wait() - ln_out_total_work = None - if fc1_dgrad_work is not None: - fc1_dgrad_work.wait() - fc1_dgrad_work = None + # -------------------------------------------------- + # Finished FC1 WGRAD... + # -------------------------------------------------- - # Residual gradient - dgrad = fc1_dgrad.view(inputmat.shape) - if ctx.return_layernorm_output and not ctx.return_layernorm_output_gathered: - dgrad = dgrad + grad_outputs[1].view_as(dgrad) + # Make sure all tensor-parallel communication is finished + if ln_out_total_work is not None: + ln_out_total_work.wait() + ln_out_total_work = None + if fc1_dgrad_work is not None: + fc1_dgrad_work.wait() + fc1_dgrad_work = None - # Norm gradient - dgamma = None + # Residual gradient + dgrad = fc1_dgrad.view(inputmat.shape) + if args.return_layernorm_output and not args.return_layernorm_output_gathered: + dgrad = dgrad + args.grad_ln_out.view_as(dgrad) + + # Norm gradient + dgamma = None + dbeta = None + if args.normalization == "LayerNorm": + dgrad, dgamma, dbeta = tex.layernorm_bwd( + dgrad, + inputmat, + mu, + rsigma, + ln_weight, + args.bwd_ln_sm_margin, + args.zero_centered_gamma, + ) + elif args.normalization == "RMSNorm": + dgrad, dgamma = tex.rmsnorm_bwd( + dgrad, + inputmat, + rsigma, + ln_weight, + args.bwd_ln_sm_margin, + args.zero_centered_gamma, + ) dbeta = None - if ctx.normalization == "LayerNorm": - dgrad, dgamma, dbeta = tex.layernorm_bwd( - dgrad, - inputmat, - mu, - rsigma, - ln_weight, - ctx.bwd_ln_sm_margin, - ctx.zero_centered_gamma, + clear_tensor_data(mu, rsigma) + + if args.fc1_weight_requires_grad: + # Handle custom DDP from mcore. + if args.fuse_wgrad_accumulation and hasattr( + fc1_weight_python_object, "grad_added_to_main_grad" + ): + fc1_weight_python_object.grad_added_to_main_grad = True + if getattr(fc1_weight_python_object, "zero_out_wgrad", False): + fc1_wgrad = torch.zeros( + fc1_weight_main_grad.shape, + dtype=fc1_weight_python_object.dtype, + device=torch.cuda.current_device(), + requires_grad=False, ) - elif ctx.normalization == "RMSNorm": - dgrad, dgamma = tex.rmsnorm_bwd( - dgrad, - inputmat, - rsigma, - ln_weight, - ctx.bwd_ln_sm_margin, - ctx.zero_centered_gamma, + else: + fc1_wgrad = torch.empty( + fc1_weight_main_grad.shape, + dtype=fc1_weight_python_object.dtype, + device=torch.cuda.current_device(), + requires_grad=False, ) - dbeta = None - clear_tensor_data(mu, rsigma) - - if ctx.fc1_weight_requires_grad: - # Handle custom DDP from mcore. - if ctx.fuse_wgrad_accumulation and hasattr( - fc1_weight_python_object, "grad_added_to_main_grad" - ): - fc1_weight_python_object.grad_added_to_main_grad = True - if getattr(fc1_weight_python_object, "zero_out_wgrad", False): - fc1_wgrad = torch.zeros( - fc1_weight_main_grad.shape, - dtype=fc1_weight_python_object.dtype, - device=torch.cuda.current_device(), - requires_grad=False, - ) - else: - fc1_wgrad = torch.empty( - fc1_weight_main_grad.shape, - dtype=fc1_weight_python_object.dtype, - device=torch.cuda.current_device(), - requires_grad=False, - ) - elif ctx.fuse_wgrad_accumulation: - fc1_wgrad = None - else: + elif args.fuse_wgrad_accumulation: fc1_wgrad = None + else: + fc1_wgrad = None - if ctx.fc2_weight_requires_grad: - # Handle custom DDP from mcore. - if ctx.fuse_wgrad_accumulation and hasattr( - fc2_weight_python_object, "grad_added_to_main_grad" - ): - fc2_weight_python_object.grad_added_to_main_grad = True - if getattr(fc2_weight_python_object, "zero_out_wgrad", False): - fc2_wgrad = torch.zeros( - fc2_weight_main_grad.shape, - dtype=fc2_weight_python_object.dtype, - device=torch.cuda.current_device(), - requires_grad=False, - ) - else: - fc2_wgrad = torch.empty( - fc2_weight_main_grad.shape, - dtype=fc2_weight_python_object.dtype, - device=torch.cuda.current_device(), - requires_grad=False, - ) - elif ctx.fuse_wgrad_accumulation: - fc2_wgrad = None - else: + if args.fc2_weight_requires_grad: + # Handle custom DDP from mcore. + if args.fuse_wgrad_accumulation and hasattr( + fc2_weight_python_object, "grad_added_to_main_grad" + ): + fc2_weight_python_object.grad_added_to_main_grad = True + if getattr(fc2_weight_python_object, "zero_out_wgrad", False): + fc2_wgrad = torch.zeros( + fc2_weight_main_grad.shape, + dtype=fc2_weight_python_object.dtype, + device=torch.cuda.current_device(), + requires_grad=False, + ) + else: + fc2_wgrad = torch.empty( + fc2_weight_main_grad.shape, + dtype=fc2_weight_python_object.dtype, + device=torch.cuda.current_device(), + requires_grad=False, + ) + elif args.fuse_wgrad_accumulation: fc2_wgrad = None + else: + fc2_wgrad = None + + # FIX THIS + # Scatter Fp8 tranposed-weight buffers + # if args.fp8: + # _fsdp_scatter_tensors( + # args.fsdp_group, + # fc1_weight_fp8 if not isinstance(fc1_weight, Float8Tensor) else None, + # fc2_weight_fp8 if not isinstance(fc2_weight, Float8Tensor) else None, + # ) + return ( + dgrad.view(args.inp_shape) if args.requires_dgrad else None, + dgamma, + dbeta, + fc1_wgrad, + fc1_bias_grad if fc1_bias is not None else None, + fc2_wgrad, # pylint: disable=possibly-used-before-assignment + fc2_bias_grad, + ) - if ctx.reduce_and_update_bwd_fp8_tensors and not is_graph_capturing(): - FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) - # FIX THIS - # Scatter Fp8 tranposed-weight buffers - # if ctx.fp8: - # _fsdp_scatter_tensors( - # ctx.fsdp_group, - # fc1_weight_fp8 if not isinstance(fc1_weight, Float8Tensor) else None, - # fc2_weight_fp8 if not isinstance(fc2_weight, Float8Tensor) else None, - # ) +class _LayerNormMLP(torch.autograd.Function): + """LayerNormMLP semi-top level module + Calls custom cuda extensions. + """ + + @staticmethod + def forward( + ctx, + inp: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: Optional[torch.Tensor], + fc1_weight: torch.Tensor, + fc1_bias: Optional[torch.Tensor], + fc2_weight: torch.Tensor, + fc2_bias: Optional[torch.Tensor], + fwd_args: LayerNormMLPFwdArgs, + ) -> Tuple[ + torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor] + ]: + """Forward pass: compute the output and set up the autograd context. + + The tensors are positional so autograd tracks them; they are + re-attached to ``fwd_args`` so every downstream helper takes a single + argument. The weight workspaces are non-differentiable cached tensors + passed in via ``fwd_args`` and the freshly produced workspaces are + returned as separate outputs so the module can refresh its cache. + """ + fwd_args.inp = inp + fwd_args.ln_weight = ln_weight + fwd_args.ln_bias = ln_bias + fwd_args.fc1_weight = fc1_weight + fwd_args.fc1_bias = fc1_bias + fwd_args.fc2_weight = fc2_weight + fwd_args.fc2_bias = fc2_bias + ( + out, + ln_out_return, + new_fc1_weight_workspace, + new_fc2_weight_workspace, + tensors_to_save_from_forward, + ctx_attrs, + ) = _layernorm_mlp_forward_impl(fwd_args) + if ctx is not None: + bwd_args = LayerNormMLPBwdArgs() + tensors_to_save_from_setup = _layernorm_mlp_setup_ctx( + bwd_args, + fwd_args, + (out, ln_out_return, new_fc1_weight_workspace, new_fc2_weight_workspace), + ctx_attrs, + tensors_to_save_from_forward, + ) + tensors_to_save, tensor_objects = prepare_for_saving(*tensors_to_save_from_setup) + ctx.save_for_backward(*tensors_to_save) + ctx.tensor_objects = tensor_objects + ctx.backward_objects = bwd_args + if not bwd_args.checkpoint and fwd_args.fp8 and fwd_args.any_requires_grad(): + bwd_args.reduce_and_update_bwd_fp8_tensors = check_fp8_reduce_and_update() + + return out, ln_out_return, new_fc1_weight_workspace, new_fc2_weight_workspace + + @staticmethod + def backward( + ctx, + grad_output: torch.Tensor, + grad_ln_out: Optional[torch.Tensor], + _grad_fc1_weight_workspace, + _grad_fc2_weight_workspace, + ) -> Tuple[Union[torch.Tensor, None], ...]: + """Backward pass: compute gradients and reduce FP8 scaling factors.""" + bwd_args: LayerNormMLPBwdArgs = ctx.backward_objects + bwd_args.grad_output = grad_output + bwd_args.grad_ln_out = grad_ln_out + with get_nvtx_range_context("_LayerNormMLP_backward"): + _layernorm_mlp_recompute(bwd_args, ctx) + ( + dgrad, + dgamma, + dbeta, + fc1_wgrad, + fc1_bias_grad, + fc2_wgrad, + fc2_bias_grad, + ) = _layernorm_mlp_backward_impl(bwd_args) + reduce_and_update_bwd_fp8_tensors = bwd_args.reduce_and_update_bwd_fp8_tensors + # Drop all references held by bwd_args (saved tensors, quantizers, weakrefs, + # main_grad closures) so they don't outlive backward via ctx under retain_graph. + ctx.backward_objects = None + del bwd_args + if reduce_and_update_bwd_fp8_tensors and not is_graph_capturing(): + FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) return ( - dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, + dgrad, dgamma, dbeta, fc1_wgrad, - None, # fc1_weight_workspace - fc1_bias_grad if fc1_bias is not None else None, - fc2_wgrad, # pylint: disable=possibly-used-before-assignment - None, # fc2_weight_workspace + fc1_bias_grad, + fc2_wgrad, fc2_bias_grad, - None, + None, # fwd_args ) @@ -2361,13 +2706,6 @@ def forward( if self.bias_gelu_nvfusion and not use_reentrant_activation_recompute(): self.fast_setattr("bias_gelu_nvfusion", False) - if is_grad_enabled: - fwd_fn = _LayerNormMLP.apply - autograd_ctx = [] - else: - fwd_fn = _LayerNormMLP.forward - autograd_ctx = [None] - cache_name_fc1 = ( None if (is_first_microbatch is None or self.is_fsdp2) else "fc1_weight" ) @@ -2381,71 +2719,151 @@ def forward( self._fp8_workspaces.get(cache_name_fc2) if cache_name_fc2 is not None else None ) - non_tensor_args = ( - self.eps, - is_first_microbatch, - self.fp8, - self.fp8_calibration, - self.wgrad_store, - self.fuse_wgrad_accumulation, - fc1_input_quantizer, - fc1_weight_quantizer, - fc1_output_quantizer, - fc1_grad_input_quantizer, - fc1_grad_weight_quantizer, - fc1_grad_output_quantizer, - fc2_input_quantizer, - fc2_weight_quantizer, - fc2_output_quantizer, - fc2_grad_input_quantizer, - fc2_grad_weight_quantizer, - fc2_grad_output_quantizer, - is_cpu_offload_enabled(), - self.tp_group, - self.tp_size, - self.sequence_parallel, - self.tp_size > 1, - self.activation_dtype, - self.return_layernorm_output, - self.return_layernorm_output_gathered, - self.bias_gelu_nvfusion and not self.fp8 and not debug, - self.set_parallel_mode, - is_grad_enabled, - self.fwd_ln_sm_margin if is_grad_enabled else self.inf_ln_sm_margin, - self.bwd_ln_sm_margin, - self.zero_centered_gamma, - self.activation, - self.activation_params, - self.normalization, - self.ub_overlap_ag, - self.ub_overlap_rs, - self.ub_overlap_rs_dgrad, - self.ub_bulk_dgrad, - self.ub_bulk_wgrad, - self.gemm_gelu_fusion and not debug, - self.fsdp_group, - self.fp8_meta, - cache_name_fc1 is not None, - skip_fp8_weight_update, - self.symmetric_ar_type, - self.checkpoint, - debug, - self.is_fsdp2, + dgrad_use_split_accumulator = _2X_ACC_DGRAD + wgrad_use_split_accumulator = _2X_ACC_WGRAD + if self.fp8: + _recipe = FP8GlobalStateManager.get_fp8_recipe() + backward_override = _recipe.backward_override + if hasattr(_recipe, "fp8_gemm_dgrad"): + dgrad_use_split_accumulator = _recipe.fp8_gemm_dgrad.use_split_accumulator + if hasattr(_recipe, "fp8_gemm_wgrad"): + wgrad_use_split_accumulator = _recipe.fp8_gemm_wgrad.use_split_accumulator + else: + backward_override = None + + if debug: # turn off userbuffers in debug mode + ub_overlap_ag = False + ub_overlap_rs = False + ub_overlap_rs_dgrad = False + ub_bulk_wgrad = False + ub_bulk_dgrad = False + else: + ub_overlap_ag = self.ub_overlap_ag + ub_overlap_rs = self.ub_overlap_rs + ub_overlap_rs_dgrad = self.ub_overlap_rs_dgrad + ub_bulk_wgrad = self.ub_bulk_wgrad + ub_bulk_dgrad = self.ub_bulk_dgrad + + fc2_bias_tensor = ( + fc2_bias if self.apply_bias and not self.gemm_bias_unfused_add else None ) - out, ln_out, new_fc1_ws, new_fc2_ws = fwd_fn( - *autograd_ctx, - inp, - self.layer_norm_weight, - self.layer_norm_bias, - fc1_weight, - fc1_weight_workspace, - fc1_bias, - fc2_weight, - fc2_weight_workspace, - fc2_bias if self.apply_bias and not self.gemm_bias_unfused_add else None, - non_tensor_args, + fwd_args = LayerNormMLPFwdArgs( + # tensors + inp=inp, + ln_weight=self.layer_norm_weight, + ln_bias=self.layer_norm_bias, + fc1_weight=fc1_weight, + fc1_bias=fc1_bias, + fc2_weight=fc2_weight, + fc2_bias=fc2_bias_tensor, + fc1_weight_workspace=fc1_weight_workspace, + fc2_weight_workspace=fc2_weight_workspace, + # requires_grad flags + input_requires_grad=inp.requires_grad, + ln_weight_requires_grad=self.layer_norm_weight.requires_grad, + ln_bias_requires_grad=( + self.layer_norm_bias.requires_grad + if self.layer_norm_bias is not None + else False + ), + fc1_weight_requires_grad=fc1_weight.requires_grad, + fc1_bias_requires_grad=fc1_bias.requires_grad if fc1_bias is not None else False, + fc2_weight_requires_grad=fc2_weight.requires_grad, + fc2_bias_requires_grad=( + fc2_bias_tensor.requires_grad if fc2_bias_tensor is not None else False + ), + # quantizers + fc1_input_quantizer=fc1_input_quantizer, + fc1_weight_quantizer=fc1_weight_quantizer, + fc1_output_quantizer=fc1_output_quantizer, + fc1_grad_input_quantizer=fc1_grad_input_quantizer, + fc1_grad_weight_quantizer=fc1_grad_weight_quantizer, + fc1_grad_output_quantizer=fc1_grad_output_quantizer, + fc2_input_quantizer=fc2_input_quantizer, + fc2_weight_quantizer=fc2_weight_quantizer, + fc2_output_quantizer=fc2_output_quantizer, + fc2_grad_input_quantizer=fc2_grad_input_quantizer, + fc2_grad_weight_quantizer=fc2_grad_weight_quantizer, + fc2_grad_output_quantizer=fc2_grad_output_quantizer, + # normalization + eps=self.eps, + normalization=self.normalization, + zero_centered_gamma=self.zero_centered_gamma, + fwd_ln_sm_margin=( + self.fwd_ln_sm_margin if is_grad_enabled else self.inf_ln_sm_margin + ), + bwd_ln_sm_margin=self.bwd_ln_sm_margin, + return_layernorm_output=self.return_layernorm_output, + return_layernorm_output_gathered=self.return_layernorm_output_gathered, + # activation + activation=self.activation, + activation_params=self.activation_params, + bias_gelu_fusion=self.bias_gelu_nvfusion and not self.fp8 and not debug, + gemm_gelu_fusion=self.gemm_gelu_fusion and not debug, + # numerical / dtype config + activation_dtype=self.activation_dtype, + fp8=self.fp8, + fp8_calibration=self.fp8_calibration, + backward_override=backward_override, + dgrad_use_split_accumulator=dgrad_use_split_accumulator, + wgrad_use_split_accumulator=wgrad_use_split_accumulator, + debug=debug, + # weight-workspace caching + is_first_microbatch=is_first_microbatch, + cache_weight=cache_name_fc1 is not None, + skip_fp8_weight_update=skip_fp8_weight_update, + # tensor / sequence parallelism + set_parallel_mode=self.set_parallel_mode, + tp_group=self.tp_group, + tp_size=self.tp_size, + tensor_parallel=self.tp_size > 1, + sequence_parallel=self.sequence_parallel, + symmetric_ar_type=self.symmetric_ar_type, + # userbuffers + ub_overlap_ag=ub_overlap_ag, + ub_overlap_rs=ub_overlap_rs, + ub_overlap_rs_dgrad=ub_overlap_rs_dgrad, + ub_bulk_dgrad=ub_bulk_dgrad, + ub_bulk_wgrad=ub_bulk_wgrad, + # FSDP + fsdp_group=self.fsdp_group, + is_fsdp2=self.is_fsdp2, + # weight-grad scheduling + fuse_wgrad_accumulation=self.fuse_wgrad_accumulation, + wgrad_store=self.wgrad_store, + # activation checkpointing + checkpoint=self.checkpoint, + fp8_meta=self.fp8_meta if self.checkpoint else None, + recompute_for_bwd=False, + # misc + cpu_offloading=is_cpu_offload_enabled(), + is_grad_enabled=is_grad_enabled, ) + if is_grad_enabled: + out, ln_out, new_fc1_ws, new_fc2_ws = _LayerNormMLP.apply( + inp, + self.layer_norm_weight, + self.layer_norm_bias, + fc1_weight, + fc1_bias, + fc2_weight, + fc2_bias_tensor, + fwd_args, + ) + else: + out, ln_out, new_fc1_ws, new_fc2_ws = _LayerNormMLP.forward( + None, + inp, + self.layer_norm_weight, + self.layer_norm_bias, + fc1_weight, + fc1_bias, + fc2_weight, + fc2_bias_tensor, + fwd_args, + ) + if new_fc1_ws is not None and cache_name_fc1 is not None: if isinstance(new_fc1_ws, torch.Tensor): new_fc1_ws = new_fc1_ws.detach() diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 94de69e9759..105ea8dd5da 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -33,6 +33,7 @@ ) from ._common import ( can_reconstruct_wgrad_input_from_original, + check_fp8_reduce_and_update, noop_cat, set_quantizer_amax_reduction_group, set_quantizer_usage_for_wgrad_all_gather, @@ -317,16 +318,6 @@ def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None: ) # pylint: disable=unbalanced-tuple-unpacking -def _check_fp8_reduce_and_update(): - """Check if this is the first FP8 module (for backward reduce-and-update).""" - qstate = FP8GlobalStateManager.quantization_state - _first_fp8_module = qstate.is_first_fp8_module - result = FP8GlobalStateManager.is_first_fp8_module() - if in_fp8_activation_recompute_phase(): - qstate.is_first_fp8_module = _first_fp8_module - return result - - def _out_leading_from_inp(leading: int, args: Union[LinearFwdArgs, LinearBwdArgs]) -> int: """Output's leading (sequence) dim from the input's: sequence parallelism gathers it (column-parallel) or scatters it (row-parallel).""" @@ -1481,8 +1472,9 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None # its persistent buffer; cuBLASMp does not, so we gather here. Route # through the same FP8-aware all-gather as the non-overlap path in # ``TransformerEngineBaseModule.grad_output_preprocess`` by passing the - # grad_output quantizer. The columnwise data needed for wgrad is then - # produced by ``update_usage(columnwise_usage=True)`` further below. + # grad_output quantizer. Per-tensor FP8 can reconstruct columnwise + # data from the gathered rowwise data; MXFP8 must instead quantize + # the original gradient columnwise to avoid double quantization. if ( bwd_args.requires_wgrad and bwd_args.ub_overlap_ag @@ -1491,6 +1483,8 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None ): if grad_output_quantizer is not None: set_quantizer_usage_for_wgrad_all_gather(grad_output_quantizer) + if isinstance(grad_output_quantizer, MXFP8Quantizer): + grad_output = grad_output_arg.reshape(-1, grad_output_arg.shape[-1]).contiguous() grad_output, _ = gather_along_first_dim( grad_output, bwd_args.tp_group, @@ -1520,7 +1514,11 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None # Prepare grad output tensor # Note: Synchronize tensor-parallel communication and # make sure required data is available - if bwd_args.ub_overlap_ag and isinstance(grad_output_quantizer, MXFP8Quantizer): + if ( + bwd_args.ub_overlap_ag + and isinstance(grad_output_quantizer, MXFP8Quantizer) + and not ub_obj_dgrad.with_cublasmp() + ): # UB does not support pipelined overlapping grad output # all-gather with wgrad GEMM. Also, we can't # convert row-scaled MXFP8 to column-scaled, so we @@ -1852,7 +1850,7 @@ def forward( or fwd_args.weight_requires_grad or fwd_args.bias_requires_grad ): - bwd_args.reduce_and_update_bwd_fp8_tensors = _check_fp8_reduce_and_update() + bwd_args.reduce_and_update_bwd_fp8_tensors = check_fp8_reduce_and_update() if fwd_args.backward_override is not None: bwd_args.reduce_and_update_bwd_fp8_tensors = False diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index ac38f7d1cf9..898310b0d61 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -40,10 +40,39 @@ "fused_atomic": int(QBHistogramMode.FUSED_ATOMIC), } _QB_BOUNDS_VALIDATED_VERSION_ATTR = "_nvte_qb_bounds_validated_version" +_QB_BOUNDS_CAPTURE_VALIDATED_VERSION_ATTR = "_nvte_qb_bounds_capture_validated_version" + + +def mark_qb_bin_bounds_validated(bin_bounds: torch.Tensor) -> None: + """Mark the current QB bounds version as valid after a trusted device-side update. + + This function validates metadata only; it does not inspect or schedule validation of tensor + values. The caller must guarantee finite FP32 bounds with ``lower < upper``. When called during + CUDA graph capture, call it immediately after dispatching the trusted in-place update on the + same stream. The Python marker itself is not replayed, so every value produced on replay must + satisfy the bounds contract for every possible input and control-flow path. + """ + if not ( + isinstance(bin_bounds, torch.Tensor) + and bin_bounds.is_cuda + and bin_bounds.is_contiguous() + and bin_bounds.dtype == torch.float32 + and bin_bounds.shape == (2,) + ): + raise ValueError("QB bin_bounds must be a contiguous FP32 CUDA tensor with shape [2]") + version = bin_bounds._version + setattr(bin_bounds, _QB_BOUNDS_VALIDATED_VERSION_ATTR, version) + with torch.cuda.device(bin_bounds.device): + capture_validated_version = version if torch.cuda.is_current_stream_capturing() else None + setattr( + bin_bounds, + _QB_BOUNDS_CAPTURE_VALIDATED_VERSION_ATTR, + capture_validated_version, + ) def _validate_qb_bin_bounds(bin_bounds: torch.Tensor) -> bool: - """Validate CUDA-resident QB bounds once per PyTorch tensor version.""" + """Validate QB bounds once per PyTorch tensor version.""" if not ( isinstance(bin_bounds, torch.Tensor) and bin_bounds.is_cuda @@ -55,19 +84,24 @@ def _validate_qb_bin_bounds(bin_bounds: torch.Tensor) -> bool: return False version = bin_bounds._version - if getattr(bin_bounds, _QB_BOUNDS_VALIDATED_VERSION_ATTR, None) == version: + validated_version = getattr(bin_bounds, _QB_BOUNDS_VALIDATED_VERSION_ATTR, None) + capture_validated_version = getattr(bin_bounds, _QB_BOUNDS_CAPTURE_VALIDATED_VERSION_ATTR, None) + if validated_version == version and capture_validated_version != version: return True with torch.cuda.device(bin_bounds.device): - if torch.cuda.is_current_stream_capturing(): + is_capturing = torch.cuda.is_current_stream_capturing() + if validated_version == version and is_capturing: + return True + if is_capturing: raise RuntimeError( - "QB bin_bounds must be validated by an eager router call before CUDA graph capture" + "QB bin_bounds current version must be validated before CUDA graph capture" ) lower, upper = bin_bounds.detach().cpu().tolist() if not (math.isfinite(lower) and math.isfinite(upper) and lower < upper): raise ValueError( f"QB bin_bounds values must be finite with lower < upper, got [{lower}, {upper}]" ) - setattr(bin_bounds, _QB_BOUNDS_VALIDATED_VERSION_ATTR, version) + mark_qb_bin_bounds_validated(bin_bounds) return True @@ -287,8 +321,13 @@ def fused_topk_with_score_function( Caller-owned int32 ``[num_experts, num_bins]`` histogram accumulated in place. qb_bin_bounds : torch.Tensor, optional FP32 CUDA tensor ``[lower, upper]`` defining uniform QB histogram bins. Values must be - finite with ``lower < upper``. Bounds are revalidated after PyTorch-tracked in-place - updates; validate once with an eager call before CUDA graph capture. + finite with ``lower < upper``. Eager calls revalidate PyTorch-tracked in-place updates. + Before CUDA graph capture, validate the current version eagerly or use + :func:`mark_qb_bin_bounds_validated` after a trusted device-side update. CUDA graph replay + does not execute Python validation or advance PyTorch tensor-version counters. Any bounds + update visible to replay, whether issued outside the graph or captured inside it, is + therefore trusted and must produce finite values with ``lower < upper`` on every replay. + Invalid replay-time bounds may silently produce an incorrect histogram. qb_histogram_mode : str, optional ``"two_kernel"`` or ``"fused_atomic"``. Must be provided with the two QB tensors.