Reuse cuBLASMp's all-gathered grad_output in the wgrad GEMM (#3629)

* [PyTorch] Reuse cuBLASMp's all-gathered grad_output for wgrad

With the cuBLASMp comm+GEMM overlap backend, the dgrad AllGather + GEMM gathers grad_output and the wgrad GEMM
gathered it a second time, while Userbuffers wgrad reads the tensor that its dgrad all-gather left in the
communication buffer. cuBLASMp 0.11.0 leaves the gathered B at the start of its workspace (a documented
cublasMpMatmul workspace property), so wgrad now reuses it the same way:

- nvte_comm_gemm_workspace() (comm_gemm.h, TE-internal) returns the comm-gemm context's workspace.
- With the cuBLASMp backend, CommOverlap/CommOverlapP2P::get_buffer() returns a view of that workspace with the
  communicator's buffer shape and dtype, like the Userbuffers buffer view.
- get_cublasmp_all_gather_output() builds the gathered tensor from it for unquantized and per-tensor FP8 data.
  Linear, LayerNormLinear, LayerNormMLP (fc2) and the fused te.ops backward use it instead of
  gather_along_first_dim; MXFP8 and other formats, a delayed wgrad and a skipped dgrad GEMM still gather.

This needs cuBLASMp 0.11.0 or newer, since older versions keep the gathered B elsewhere in the workspace; the next
commit checks the version.

Tested on 4x H100 with cuBLASMp 0.11.0: the layer and multi-layer tests of test_comm_gemm_overlap.py with the
cuBLASMp backend pass (39 passed, 13 MXFP8 cases skipped on H100), and test_fusible_ops_with_userbuffers.py
--use-cublasmp passes in BF16 and FP8 delayed and current scaling. All-gathers per backward: row-parallel Linear
1 -> 0, LayerNormMLP 2 -> 1 (fc1 still re-gathers its input, as with Userbuffers).

Signed-off-by: Almog Segal <asegal@nvidia.com>

* [PyTorch] Fall back to the all-gather with cuBLASMp before 0.11.0

Reusing the gathered grad_output relies on cuBLASMp 0.11.0 leaving the gathered B at the start of its workspace.
Earlier versions keep it after their cuBLASLt workspaces, so the view held other data. nvte_comm_gemm_workspace() now
checks the loaded library with cublasMpGetVersion() and returns nullptr before 0.11.0; get_buffer() then reports that
no gathered input is available, and get_cublasmp_all_gather_output() returns None, so the callers all-gather as before.

Tested on 4x H100: with cuBLASMp 0.11.0, the row-parallel Linear backward issues no all-gather and the Linear and
LayerNormMLP overlap tests with the cuBLASMp backend pass (9 passed). With cuBLASMp 0.10.0 it all-gathers once per
backward and the same tests pass, as they do on TE main.

Signed-off-by: Almog Segal <asegal@nvidia.com>

* [PyTorch] Check the cuBLASMp version instead of catching the get_buffer() error

get_cublasmp_all_gather_output() fell back to the all-gather for cuBLASMp versions before 0.11.0 by catching the
error that get_buffer() raises and matching its message. That breaks silently if the message changes, and with an
older cuBLASMp it raises and catches an exception in every backward pass of every layer that overlaps an all-gather.

nvte_cublasmp_version() now returns the version of the loaded cuBLASMp library, read once with cublasMpGetVersion()
(0 without cuBLASMp), and Python exposes it as tex.nvte_cublasmp_version(). get_cublasmp_all_gather_output() checks a
cached version >= 0.11.0 before calling get_buffer(), so get_buffer() errors are no longer caught.
nvte_comm_gemm_workspace() uses the same function.

Signed-off-by: Almog Segal <asegal@nvidia.com>

* [PyTorch] Test the fused te.ops comm+GEMM overlap with cuBLASMp

test_fuser_ops_with_userbuffers only ran the Userbuffers backend, although the script supports --use-cublasmp and the
fused te.ops backward reuses cuBLASMp's gathered grad_output. It is now parametrized over both backends and skips the
cuBLASMp cases when Transformer Engine is built without cuBLASMp.

Signed-off-by: Almog Segal <asegal@nvidia.com>

---------

Signed-off-by: Almog Segal <asegal@nvidia.com>

Author: almogsegal

tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py

@@ -431,14 +431,19 @@ def run_parallel_tests(model_config: ModelConfig) -> None:
 
 @pytest.mark.parametrize("world_size", _world_sizes)
 @pytest.mark.parametrize("quantization", quantization_list)
+@pytest.mark.parametrize("use_cublasmp", (False, True), ids=("userbuffers", "cublasmp"))
 def test_fuser_ops_with_userbuffers(
     *,
     world_size: int,
     dtype: torch.dtype = torch.bfloat16,
     quantization: Optional[str],
+    use_cublasmp: bool,
 ) -> None:
     """Launch parallel job and run tests"""
 
+    if use_cublasmp and not tex.nvte_built_with_cublasmp():
+        pytest.skip("Transformer Engine not built with cuBLASMp (NVTE_WITH_CUBLASMP=0).")
+
     # Parallel job launcher
     command = []
     if tex.ubuf_built_with_mpi():
@@ -462,6 +467,8 @@ def test_fuser_ops_with_userbuffers(
     )
     if quantization is not None:
         command.extend(("--quantization", quantization))
+    if use_cublasmp:
+        command.append("--use-cublasmp")
 
     # Environment
     env = dict(os.environ)

transformer_engine/common/comm_gemm/comm_gemm.cpp

@@ -602,6 +602,26 @@ int64_t nvte_comm_gemm_numroc(NVTECommGemmCtx* ctx, int64_t global_size) {
   return cublasMpNumroc(global_size, block_size(ctx, global_size), ctx->rank, 0, ctx->nranks);
 }
 
+void* nvte_comm_gemm_workspace(NVTECommGemmCtx* ctx, size_t* size) {
+  NVTE_API_CALL(nvte_comm_gemm_workspace);
+  // The loaded library decides the workspace layout: cuBLASMp 0.11.0 and newer start it with the
+  // gathered B.
+  if (nvte_cublasmp_version() < 1100) {
+    *size = 0;
+    return nullptr;
+  }
+  *size = ctx->workspace_size;
+  return ctx->workspace;
+}
+
+int nvte_cublasmp_version() {
+  static const int version = [] {
+    int v = 0;
+    return cublasMpGetVersion(&v) == CUBLASMP_STATUS_SUCCESS ? v : 0;
+  }();
+  return version;
+}
+
 #else  // NVTE_WITH_CUBLASMP
 
 struct NVTECommGemmCtx {};
@@ -642,4 +662,10 @@ int64_t nvte_comm_gemm_numroc(NVTECommGemmCtx* ctx, int64_t global_size) {
   NVTE_ERROR("Transformer Engine has not been built with cuBLASMp support.");
 }
 
+void* nvte_comm_gemm_workspace(NVTECommGemmCtx* ctx, size_t* size) {
+  NVTE_ERROR("Transformer Engine has not been built with cuBLASMp support.");
+}
+
+int nvte_cublasmp_version() { return 0; }
+
 #endif  // NVTE_WITH_CUBLASMP

transformer_engine/common/include/transformer_engine/comm_gemm.h

@@ -151,6 +151,27 @@ void nvte_gemm_all_reduce(NVTECommGemmCtx* ctx, int64_t m, int64_t n, int64_t k,
  */
 int64_t nvte_comm_gemm_numroc(NVTECommGemmCtx* ctx, int64_t global_size);
 
+/*! \brief Get the device workspace of a comm-gemm context.
+ *
+ *  After nvte_all_gather_gemm with transb = false completes on its stream, the workspace begins
+ *  with the gathered B: the local B of every rank, in rank order, which together form the
+ *  gathered (N, K) tensor. The data remains valid until the next operation with the context on
+ *  any rank, since the other ranks write into the workspace. cuBLASMp versions before 0.11.0 keep
+ *  the gathered B elsewhere, so with those this function returns nullptr.
+ *
+ *  \param[in]  ctx   Comm-GEMM context.
+ *  \param[out] size  Workspace size in bytes (0 when the function returns nullptr).
+ *  \return Workspace pointer, or nullptr before the first operation or with cuBLASMp before 0.11.0.
+ */
+void* nvte_comm_gemm_workspace(NVTECommGemmCtx* ctx, size_t* size);
+
+/*! \brief Get the version of the cuBLASMp library loaded at run time.
+ *
+ *  \return MAJOR * 10000 + MINOR * 100 + PATCH (for example, 1100 for 0.11.0), or 0 if
+ *          Transformer Engine has not been built with cuBLASMp support.
+ */
+int nvte_cublasmp_version(void);
+
 #ifdef __cplusplus
 }  // extern "C"
 #endif

transformer_engine/common/util/pybind_helper.h

@@ -163,6 +163,7 @@
   m.def("ubuf_built_with_mpi", &transformer_engine::ubuf_built_with_mpi,                           \
         py::call_guard<py::gil_scoped_release>());                                                 \
   m.def("nvte_built_with_cublasmp", &nvte_built_with_cublasmp,                                     \
-        py::call_guard<py::gil_scoped_release>());
+        py::call_guard<py::gil_scoped_release>());                                                 \
+  m.def("nvte_cublasmp_version", &nvte_cublasmp_version, py::call_guard<py::gil_scoped_release>());
 
 #endif

transformer_engine/pytorch/csrc/extensions.h

@@ -878,6 +878,9 @@ class CommOverlap : torch::CustomClassHolder, public transformer_engine::CommOve
   // Keeps the cuBLASMp NCCL communicator alive for the lifetime of this
   // instance, independent of the CommOverlapHelper that created it.
   CommOverlapHelper::NcclCommSharedPtr _nccl_comm;
+  // cuBLASMp: shape and dtype of the buffer view returned by get_buffer()
+  std::vector<size_t> _cublasmp_buffer_shape;
+  at::ScalarType _cublasmp_buffer_dtype{at::kByte};
 
  public:
   CommOverlap(const std::vector<size_t> &buffer_shape, at::ScalarType buffer_dtype,
@@ -891,6 +894,8 @@ class CommOverlap : torch::CustomClassHolder, public transformer_engine::CommOve
   // the construction-time warmup matmul that primes cuBLASMp's lazy NCCL
   // window registrations and workspace allocation so subsequent matmuls
   // (including those captured in CUDA graphs) avoid the unsafe lazy paths.
+  // get_buffer() returns a view of cuBLASMp's workspace with `buffer_shape`
+  // and `buffer_dtype`: after an all-gather GEMM, it holds the gathered input.
   CommOverlap(CommOverlapHelper *helper, int tp_rank, int tp_size,
               transformer_engine::CommOverlapType comm_type,
               const std::vector<size_t> &buffer_shape, at::ScalarType buffer_dtype,
@@ -913,6 +918,9 @@ class CommOverlapP2P : torch::CustomClassHolder, public transformer_engine::Comm
   // Keeps the cuBLASMp NCCL communicator alive for the lifetime of this
   // instance, independent of the CommOverlapHelper that created it.
   CommOverlapHelper::NcclCommSharedPtr _nccl_comm;
+  // cuBLASMp: shape and dtype of the buffer view returned by get_buffer()
+  std::vector<size_t> _cublasmp_buffer_shape;
+  at::ScalarType _cublasmp_buffer_dtype{at::kByte};
 
  public:
   CommOverlapP2P(const std::vector<size_t> &buffer_shape, at::ScalarType buffer_dtype,

transformer_engine/pytorch/csrc/extensions/comm_gemm_overlap.cpp

@@ -307,17 +307,44 @@ void cublasmp_capture_warmup(te::CommOverlapCore *core, int tp_size, te::CommOve
   cudaFree(d_ptr);
 }
 
+// cuBLASMp has no communication buffer: its all-gather GEMM leaves the gathered input at the
+// start of its workspace, the local input of every rank in rank order. Like a Userbuffers
+// buffer, the view is valid until the next communicator operation on any rank overwrites it.
+at::Tensor cublasmp_get_buffer(NVTECommGemmCtx *ctx, int tp_rank, int tp_size,
+                               const std::vector<size_t> &buffer_shape, at::ScalarType dtype,
+                               bool local_chunk, std::optional<std::vector<int64_t>> shape) {
+  if (!shape) {
+    NVTE_CHECK(buffer_shape.size() == 2, "cuBLASMp buffer shape must be 2-dimensional");
+    int64_t dim0 = static_cast<int64_t>(buffer_shape[0]);
+    if (local_chunk) {
+      dim0 /= tp_size;
+    }
+    shape = {dim0, static_cast<int64_t>(buffer_shape[1])};
+  }
+  const size_t bytes = transformer_engine::pytorch::product(*shape) * at::elementSize(dtype);
+  const size_t offset = local_chunk ? bytes * tp_rank : 0;
+  size_t workspace_size = 0;
+  char *workspace = static_cast<char *>(nvte_comm_gemm_workspace(ctx, &workspace_size));
+  NVTE_CHECK(workspace != nullptr,
+             "No gathered input in the cuBLASMp workspace (requires cuBLASMp 0.11.0 or newer and a "
+             "previous all-gather GEMM)");
+  NVTE_CHECK(offset + bytes <= workspace_size, "Requested buffer (shape=", *shape,
+             ", local_chunk=", local_chunk, ") exceeds the cuBLASMp workspace (", workspace_size,
+             " bytes)");
+  return torch::from_blob(workspace + offset, *shape, at::dtype(dtype).device(torch::kCUDA));
+}
+
 }  // namespace
 
 CommOverlap::CommOverlap(CommOverlapHelper *helper, int tp_rank, int tp_size,
                          te::CommOverlapType comm_type, const std::vector<size_t> &buffer_shape,
                          at::ScalarType buffer_dtype, int num_comm_sm, bool atomic_gemm)
     : te::CommOverlapBase(helper->get_nccl_comm("intra").get(), tp_rank, tp_size, num_comm_sm,
                           atomic_gemm),
-      _nccl_comm(helper->get_nccl_comm("intra")) {
-  // buffer_dtype is unused on this path (the warmup runs in BF16); kept in
-  // the signature for API symmetry with the non-cuBLASMp ctor.
-  (void)buffer_dtype;
+      _nccl_comm(helper->get_nccl_comm("intra")),
+      _cublasmp_buffer_shape(buffer_shape),
+      _cublasmp_buffer_dtype(buffer_dtype) {
+  // The warmup runs in BF16; buffer_dtype is the dtype of get_buffer() views.
   cublasmp_capture_warmup(this, tp_size, comm_type, buffer_shape);
 }
 
@@ -362,6 +389,11 @@ void CommOverlap::copy_into_buffer(const at::Tensor &input, bool local_chunk) {
 }
 
 at::Tensor CommOverlap::get_buffer(bool local_chunk, std::optional<std::vector<int64_t>> shape) {
+  if (_with_cublasmp) {
+    return cublasmp_get_buffer(_cublasmp_ctx, _tp_id, _tp_size, _cublasmp_buffer_shape,
+                               _cublasmp_buffer_dtype, local_chunk, shape);
+  }
+
   // Check buffer shape
   const size_t ubuf_size = _ubuf.numel();
   if (shape) {
@@ -426,9 +458,10 @@ CommOverlapP2P::CommOverlapP2P(CommOverlapHelper *helper, int tp_rank, int tp_si
                                int num_comm_sm, bool atomic_gemm)
     : te::CommOverlapP2PBase(helper->get_nccl_comm("intra").get(), tp_rank, tp_size, num_comm_sm,
                              atomic_gemm),
-      _nccl_comm(helper->get_nccl_comm("intra")) {
-  // See CommOverlap constructor for the buffer_dtype rationale.
-  (void)buffer_dtype;
+      _nccl_comm(helper->get_nccl_comm("intra")),
+      _cublasmp_buffer_shape(buffer_shape),
+      _cublasmp_buffer_dtype(buffer_dtype) {
+  // The warmup runs in BF16; buffer_dtype is the dtype of get_buffer() views.
   cublasmp_capture_warmup(this, tp_size, comm_type, buffer_shape);
 }
 
@@ -470,6 +503,11 @@ void CommOverlapP2P::copy_into_buffer(const at::Tensor &input, bool local_chunk)
 }
 
 at::Tensor CommOverlapP2P::get_buffer(bool local_chunk, std::optional<std::vector<int64_t>> shape) {
+  if (_with_cublasmp) {
+    return cublasmp_get_buffer(_cublasmp_ctx, _tp_id, _tp_size, _cublasmp_buffer_shape,
+                               _cublasmp_buffer_dtype, local_chunk, shape);
+  }
+
   // Check buffer shape
   if (shape) {
     const size_t requested_size = transformer_engine::pytorch::product(*shape);

transformer_engine/pytorch/module/base.py

@@ -3,6 +3,7 @@
 # See LICENSE for license information.
 
 """Base modules and utilities for TransformerEngine PyTorch API"""
+import functools
 import io
 import math
 import os
@@ -812,6 +813,63 @@ def fill_userbuffers_buffer_for_all_gather(
     raise ValueError(f"Unsupported quantizer for Userbuffers ({quantizer})")
 
 
+@functools.lru_cache(maxsize=None)
+def _cublasmp_keeps_gathered_input() -> bool:
+    """Whether the loaded cuBLASMp leaves the gathered all-gather GEMM input at the start of its workspace
+
+    cuBLASMp 0.11.0 and newer do; earlier versions keep it elsewhere.
+
+    """
+    return tex.nvte_cublasmp_version() >= 1100
+
+
+def get_cublasmp_all_gather_output(
+    comm,
+    local_tensor: torch.Tensor | QuantizedTensorStorage,
+    quantizer: Optional[Quantizer],
+    process_group,
+) -> Optional[torch.Tensor | QuantizedTensorStorage]:
+    """Get the tensor gathered by the last cuBLASMp all-gather GEMM
+
+    cuBLASMp leaves the gathered input at the start of its workspace,
+    which ``comm.get_buffer`` views. As with the Userbuffers buffer of
+    ``fill_userbuffers_buffer_for_all_gather``, the returned tensor
+    uses that memory and is overwritten by the next operation with
+    ``comm`` on any rank. ``local_tensor`` must be the input of that
+    GEMM.
+
+    Returns ``None`` for formats that wgrad cannot use as gathered
+    (anything but unquantized and per-tensor FP8 data) and with
+    cuBLASMp versions before 0.11.0, which keep the gathered input
+    elsewhere; callers all-gather those separately.
+
+    """
+    if not _cublasmp_keeps_gathered_input():
+        return None
+    if isinstance(local_tensor, Float8TensorStorage):
+        local_data = local_tensor._data
+    elif isinstance(local_tensor, QuantizedTensorStorage):
+        return None
+    else:
+        local_data = local_tensor
+    if local_data is None:
+        return None
+    global_shape = list(local_data.size())
+    global_shape[0] *= torch.distributed.get_world_size(process_group)
+    global_data = comm.get_buffer(shape=global_shape)
+    if global_data.dtype != local_data.dtype:
+        return None
+    if not isinstance(local_tensor, Float8TensorStorage):
+        return global_data
+    return Float8TensorStorage(
+        data=global_data,
+        fp8_scale_inv=local_tensor._scale_inv,
+        fp8_dtype=local_tensor._fp8_dtype,
+        fake_dtype=local_tensor._dtype,
+        quantizer=quantizer,
+    )
+
+
 def _is_weight_workspace_valid(
     workspace: QuantizedTensorStorage,
     quantizer: Quantizer,

transformer_engine/pytorch/module/layernorm_linear.py

@@ -19,6 +19,7 @@
 from transformer_engine.pytorch.tensor.utils import clear_columnwise_cache, is_custom
 from .base import (
     fill_userbuffers_buffer_for_all_gather,
+    get_cublasmp_all_gather_output,
     get_ub,
     get_ub_is_fp8,
     is_ub_initialized,
@@ -1182,12 +1183,13 @@ def _layernorm_linear_backward_impl(
         # 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
+        # cuBLASMp's AG+GEMM leaves the gathered grad_output at the start of
+        # its workspace, as Userbuffers leaves it in its persistent buffer,
+        # so wgrad reuses it for unquantized and per-tensor FP8 data. A
+        # delayed wgrad runs after the workspace is reused, and other formats
+        # need data the dgrad GEMM did not gather, so those 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 (
@@ -1196,13 +1198,23 @@ def _layernorm_linear_backward_impl(
             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,
-            )
+            grad_output_total = None
+            if not (args.wgrad_store is not None and args.wgrad_store.delay_wgrad_compute()):
+                grad_output_total = get_cublasmp_all_gather_output(
+                    args.ub_obj_gradout,
+                    grad_output,
+                    args.grad_output_quantizer,
+                    args.tp_group,
+                )
+            if grad_output_total is None:
+                if args.grad_output_quantizer is not None:
+                    set_quantizer_usage_for_wgrad_all_gather(args.grad_output_quantizer)
+                grad_output_total, _ = gather_along_first_dim(
+                    grad_output,
+                    args.tp_group,
+                    quantizer=args.grad_output_quantizer,
+                )
+            grad_output = grad_output_total
 
         # --------------------------------------------------
         # Compute grad weight