Bulk AG Overlap for bf16 on gfx950

tests/pytorch/distributed/run_gemm_with_overlap.py

@@ -19,6 +19,7 @@
 import torch
 import torch.distributed as dist
 from torch.distributed.elastic.multiprocessing.errors import record
+from torch.utils.cpp_extension import IS_HIP_EXTENSION
 
 import transformer_engine.pytorch as te
 from transformer_engine.pytorch import (
@@ -186,8 +187,18 @@ def _parse_args(argv=None, namespace=None):
     )
     opts = parser.parse_args(argv, namespace)
 
+    if opts.fused and not IS_HIP_EXTENSION:
+        warnings.warn("The fused AG+GEMM backend is ROCm only.")
+        opts.fused = False
+
     if opts.bulk_overlap:
-        if opts.p2p:
+        if opts.fused and opts.comm_type != tex.CommOverlapType.AG:
+            warnings.warn("The fused bulk overlap is all-gather only.")
+            opts.fused = False
+        if opts.fused:
+            # `fused_overlap_bulk_ag` is a CommOverlapP2P entry point
+            opts.p2p = True
+        elif opts.p2p:
             warnings.warn("Point-2-point comms are not supported with bulk overlap.")
             opts.p2p = False
         if opts.atomic:
@@ -419,6 +430,8 @@ def dist_print(msg, src=None, info=False, error=False, section=False, group=None
         # Bulk overlap weight and input tensors are not relevant so they're globally sized
         local_kernel_t_shape = (ffn_hidden_size, hidden_size)
         local_inp_shape = (outer_size, hidden_size)
+        if opts.fused:
+            local_inp_shape = (outer_size, ffn_hidden_size)
         # Bulk overlap comm tensor is distributed for AG overlap only
         if opts.comm_type == tex.CommOverlapType.AG:
             bulk_inp_shape = (outer_size // tp_size, hidden_size)
@@ -709,11 +722,15 @@ def _fp8_gemm2(gemm1_out):
             extra_output=rs_out2,
         )
 
+    # The fused bulk all-gather GEMM is the NN one shaped above.
+    gemm_layout = "NN" if (opts.bulk_overlap and opts.fused) else "TN"
+
     def _gemm():
         return tex.general_gemm(
             kernel_t,
             gemm_inp,
             out_dtype=torch.bfloat16,
+            layout=gemm_layout,
             use_split_accumulator=te.module.base._2X_ACC_FPROP,
             ub=ub_obj,
             ub_type=opts.comm_type,

tests/pytorch/distributed/run_layer_with_overlap.py

@@ -479,6 +479,9 @@ def dist_print(msg, src=None, end="\n", debug=False, error=False):
         with_cublasmp=opts.use_cublasmp,
     )
 
+    dist_print("UB FUSED NAMES: " + " ".join(sorted(te.module.base._ub_fused_names)))
+    dist_print("UB DISABLED NAMES: " + " ".join(sorted(te.module.base._ub_disabled_names)))
+
     with te.quantized_model_init(enabled=opts.fp8_init):
         test_model = multi_module_model(opts.layer_type, opts.num_layers, *args, **kwargs)
     dist_print("Initialized test model...", debug=True)
@@ -561,6 +564,10 @@ def run_fwd_bwd(model, x):
             del test_graph
     else:
         test_out = run_fwd_bwd(test_model, test_x)
+    dist_print(
+        "UB BULK ELIGIBLE: "
+        + " ".join(sorted(n for n, ok in te.module.base._ub_fused_bulk_decisions.items() if ok))
+    )
     test_grads = [test_out, test_x.grad]
     names = ["output", "input.grad"]
     for test_name, test_param in test_model.named_parameters():

tests/pytorch/distributed/test_comm_gemm_overlap.py

@@ -120,9 +120,6 @@ def _run_layer_with_overlap(
     num_layers=1,
     use_cublasmp=False,
 ):
-    # Skip BULK overlap tests on HIP (column parallel or None with overlap_rs_dgrad=False)
-    if IS_HIP_EXTENSION and not overlap_rs_dgrad and linear_parallel_mode in ("column", None):
-        pytest.skip("Bulk overlap is not yet supported on HIP/ROCm.")
     # Reduce-scatter comm+GEMM overlap is flaky on gfx950
     if IS_HIP_EXTENSION and (overlap_rs_dgrad or linear_parallel_mode == "row"):
         pytest.skip("Reduce-scatter comm+GEMM overlap is flaky on gfx950")
@@ -480,9 +477,9 @@ def _fused_launch_cmd(nprocs: int):
     return ["torchrun", f"--nproc_per_node={nprocs}"]
 
 
-def _run_fused_ag(quantization="none", nprocs=None):
+def _run_fused_ag(nprocs, bulk=False, quantization="none"):
     """Run the AG overlap harness with the fused backend, returning the completed process."""
-    test_cmd = _fused_launch_cmd(nprocs if nprocs is not None else FUSED_PROC_COUNTS[0]) + [
+    test_cmd = _fused_launch_cmd(nprocs) + [
         str(TEST_ROOT / "run_gemm_with_overlap.py"),
         "--check-numerics",
         f"--seed={RNG_SEED}",
@@ -491,12 +488,47 @@ def _run_fused_ag(quantization="none", nprocs=None):
         f"--num-heads={NUM_HEADS}",
         f"--head-dim={HEAD_DIM}",
         "--comm-type=AG",
-        "--p2p",
         "--fused",
-        f"--quantization={quantization}",
     ]
+    test_cmd += ["--bulk-overlap"] if bulk else ["--p2p", f"--quantization={quantization}"]
     return subprocess.run(test_cmd, env=os.environ, capture_output=True, check=False)
 
+ELIGIBLE_OUT_FEATURES_PER_RANK = 1536
+INELIGIBLE_OUT_FEATURES_PER_RANK = 1568
+UNALIGNED_SEQ_LENGTH = 1152
+
+def _run_fused_layer(nprocs, extra_args, seq_length=SEQ_LENGTH):
+    """Run the layer harness on a column-parallel LayerNormLinear with the fused backend live."""
+    test_cmd = (
+        _fused_launch_cmd(nprocs)
+        + [
+            str(TEST_ROOT / "run_layer_with_overlap.py"),
+            f"--seed={RNG_SEED}",
+            f"--seq-length={seq_length}",
+            f"--batch-size={BATCH_SIZE}",
+            f"--num-heads={NUM_HEADS}",
+            f"--head-dim={HEAD_DIM}",
+            f"--layer-type={te.LayerNormLinear.__name__}",
+            "--linear-parallel-mode=column",
+            "--num-layers=1",
+            "--use-bf16-params",
+        ]
+        + extra_args
+    )
+    env = os.environ.copy()
+    env["PYTORCH_JIT"] = "0"
+    env["NVTE_TORCH_COMPILE"] = "0"
+    env["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "0"
+    return subprocess.run(test_cmd, env=env, capture_output=True, check=False)
+
+
+def _reported_names(stdout, prefix):
+    """The layer name sets the harness printed under `prefix`."""
+    for line in stdout.decode().splitlines():
+        if prefix in line:
+            return set(line.split(prefix, 1)[1].split())
+    return None
+
 
 def _assert_numerics_passed(result):
     stdout, stderr = result.stdout.decode(), result.stderr.decode()
@@ -509,7 +541,7 @@ def _assert_numerics_passed(result):
 @pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS)
 def test_fused_ag_overlap_bf16(nprocs):
     """bf16 at an aligned shape: the fused backend runs and the result is correct."""
-    _assert_numerics_passed(_run_fused_ag(nprocs=nprocs))
+    _assert_numerics_passed(_run_fused_ag(nprocs))
 
 
 @pytest.mark.skipif(not fused_available, reason=reason_for_no_fused)
@@ -521,7 +553,7 @@ def test_fused_ag_overlap_rejects_non_bf16(quantization, nprocs):
         pytest.skip(reason_for_no_fp8)
     if quantization == "mxfp8" and not mxfp8_available:
         pytest.skip(reason_for_no_mxfp8)
-    result = _run_fused_ag(quantization=quantization, nprocs=nprocs)
+    result = _run_fused_ag(nprocs, quantization=quantization)
     assert result.returncode != 0, "fused AG+GEMM accepted a non-bf16 operand"
     assert "non-bf16 operand" in result.stderr.decode(), result.stderr.decode()
 
@@ -530,9 +562,9 @@ def test_fused_ag_overlap_rejects_non_bf16(quantization, nprocs):
 @pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS)
 def test_fused_ag_overlap_is_deterministic(nprocs):
     """Bitwise reproducibility across runs"""
-    first = _run_fused_ag(nprocs=nprocs)
+    first = _run_fused_ag(nprocs)
     _assert_numerics_passed(first)
-    second = _run_fused_ag(nprocs=nprocs)
+    second = _run_fused_ag(nprocs)
     _assert_numerics_passed(second)
 
     def _hashes(out):
@@ -542,3 +574,62 @@ def _hashes(out):
     first_hashes, second_hashes = _hashes(first.stdout), _hashes(second.stdout)
     assert first_hashes, f"harness printed no output hash\n{first.stdout.decode()}"
     assert first_hashes == second_hashes, "two identical runs produced different outputs"
+
+
+@pytest.mark.skipif(not fused_available, reason=reason_for_no_fused)
+@pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS)
+def test_fused_bulk_ag_overlap_bf16(nprocs):
+    """The bulk all-gather that rides in an unrelated GEMM's grid."""
+    _assert_numerics_passed(_run_fused_ag(nprocs, bulk=True))
+
+
+@pytest.mark.skipif(not fused_available, reason=reason_for_no_fused)
+@pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS)
+def test_fused_layer_bulk_dgrad_bf16(nprocs):
+    """A column-parallel layer whose dgrad dimensions clear the fused contract."""
+    result = _run_fused_layer(nprocs, [f"--out-features={ELIGIBLE_OUT_FEATURES_PER_RANK * nprocs}"])
+    _assert_numerics_passed(result)
+    fused = _reported_names(result.stdout, "UB FUSED NAMES: ")
+    assert fused is not None, f"harness printed no fused name set\n{result.stdout.decode()}"
+    assert "qkv_dgrad" in fused, fused
+    eligible = _reported_names(result.stdout, "UB BULK ELIGIBLE: ")
+    assert eligible is not None, f"harness printed no eligibility set\n{result.stdout.decode()}"
+    assert "qkv_dgrad" in eligible, eligible
+
+
+@pytest.mark.skipif(not fused_available, reason=reason_for_no_fused)
+@pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS)
+def test_fused_layer_declines_ineligible_k(nprocs):
+    """A K the fused kernels cannot serve has to fall back to no overlap."""
+    result = _run_fused_layer(
+        nprocs, [f"--out-features={INELIGIBLE_OUT_FEATURES_PER_RANK * nprocs}"]
+    )
+    _assert_numerics_passed(result)
+    stderr = result.stderr.decode()
+    assert "ineligible shape" not in stderr, stderr
+    assert "failed to launch" not in stderr, stderr
+    fused = _reported_names(result.stdout, "UB FUSED NAMES: ")
+    disabled = _reported_names(result.stdout, "UB DISABLED NAMES: ")
+    assert fused is not None, f"harness printed no fused name set\n{result.stdout.decode()}"
+    assert "qkv_dgrad" in fused, fused
+    assert disabled is not None and "qkv_dgrad" not in disabled, disabled
+    eligible = _reported_names(result.stdout, "UB BULK ELIGIBLE: ")
+    assert eligible is not None, f"harness printed no eligibility set\n{result.stdout.decode()}"
+    assert "qkv_dgrad" not in eligible, eligible
+
+
+@pytest.mark.skipif(not fused_available, reason=reason_for_no_fused)
+@pytest.mark.parametrize("nprocs", FUSED_PROC_COUNTS)
+def test_fused_layer_declines_unaligned_region(nprocs):
+    """A Userbuffers region the fused backend cannot serve declines at setup."""
+    result = _run_fused_layer(
+        nprocs,
+        [f"--out-features={ELIGIBLE_OUT_FEATURES_PER_RANK * nprocs}"],
+        seq_length=UNALIGNED_SEQ_LENGTH,
+    )
+    _assert_numerics_passed(result)
+    fused = _reported_names(result.stdout, "UB FUSED NAMES: ")
+    disabled = _reported_names(result.stdout, "UB DISABLED NAMES: ")
+    assert fused == set(), f"expected no fused communicators, got {fused}"
+    assert disabled is not None, f"harness printed no disabled name set\n{result.stdout.decode()}"
+    assert {"qkv_fprop", "qkv_dgrad", "qkv_wgrad"} <= disabled, disabled

transformer_engine/common/comm_gemm_overlap/rocm_comm_gemm_overlap.cpp

@@ -281,6 +281,44 @@ static bool hk_fused_ag_gemm(const TensorWrapper &A, bool transa, bool transb, T
       tp_id, tp_size, chunk.bytes(), workspace.dptr(), workspace.bytes(), stream};
   return kittens_fused_ag_gemm_bf16(args);
 }
+
+// Bulk sibling of hk_fused_ag_gemm. AG is not associated with the GEMM.
+static bool hk_bulk_ag_gemm(const TensorWrapper &A, bool transa, const TensorWrapper &B, bool transb,
+                            TensorWrapper &D, const TensorWrapper &bias,
+                            const TensorWrapper &pre_gelu_out, TensorWrapper &workspace,
+                            bool accumulate, const TensorWrapper &ubuf, const TensorWrapper &chunk,
+                            communicator *comm, int reg, int tp_id, int tp_size, uint64_t signal,
+                            cudaStream_t stream) {
+  NVTE_CHECK(!transa, "fused bulk AG is NN only");
+  NVTE_CHECK(!transb && !accumulate && bias.numel() == 0 && pre_gelu_out.numel() == 0,
+             "fused bulk AG reached with an unsupported epilogue");
+  NVTE_CHECK(A.dtype() == DType::kBFloat16 && B.dtype() == DType::kBFloat16 &&
+                 D.dtype() == DType::kBFloat16 && ubuf.dtype() == DType::kBFloat16,
+             "fused bulk AG reached with a non-bf16 operand");
+  NVTE_CHECK(ubuf.numel() != 0, "fused bulk AG reached without a gather destination");
+
+  const size_t m       = A.size(1);
+  const size_t k       = A.size(0);
+  const size_t n_chunk = chunk.size(0);
+  NVTE_CHECK((tp_size == 4 || tp_size == 8) && m % 256 == 0 && k % 128 == 0 && k >= 256 && n_chunk % 256 == 0,
+             "fused bulk AG reached with an ineligible shape (m=", m, " k=", k, " n_chunk=", n_chunk,
+             " tp_size=", tp_size, ")");
+  NVTE_CHECK(D.size(0) == n_chunk * tp_size,
+             "fused bulk AG: the GEMM writes ", D.size(0), " rows but the kernel grid is sized for ",
+             n_chunk * tp_size, " from the Userbuffers region.");
+
+  const int rank_round_tp = comm->myrank - tp_id;
+  KittensFusedAgGemmArgs args{
+      A.dptr(), B.dptr(), D.dptr(),
+      reinterpret_cast<char *>(comm->gpu_ptrs) + reg * comm->nvsize * sizeof(void *),
+      rank_round_tp % comm->nvsize, comm->nvsize,
+      GET_RECV_PTR_BY_INDEX(rank_round_tp, comm, reg, 0), comm->gpu_ptrs,
+      static_cast<size_t>(GET_SEND_PTR_BY_INDEX(0, comm, reg, 0) - reinterpret_cast<char *>(comm->peer_ptr[0][0])),
+      static_cast<size_t>(GET_RECV_PTR_BY_INDEX(1, comm, reg, 0) - GET_RECV_PTR_BY_INDEX(0, comm, reg, 0)),
+      signal, static_cast<int>(m), static_cast<int>(n_chunk * tp_size), static_cast<int>(k), transa,
+      tp_id, tp_size, chunk.bytes(), workspace.dptr(), workspace.bytes(), stream, ubuf.dptr()};
+  return kittens_bulk_ag_gemm_bf16(args);
+}
 #endif
 
 void CommOverlapP2PBase::fused_overlap_ag(const TensorWrapper &A, bool transa, const TensorWrapper &B,
@@ -302,6 +340,24 @@ void CommOverlapP2PBase::fused_overlap_ag(const TensorWrapper &A, bool transa, c
   NVTE_ERROR("fused AG+GEMM was selected but is not built into this library");
 }
 
+void CommOverlapP2PBase::fused_overlap_bulk_ag(const TensorWrapper &A, bool transa,
+                                const TensorWrapper &B, bool transb, TensorWrapper &D,
+                                TensorWrapper &bias, TensorWrapper &pre_gelu_out,
+                                TensorWrapper &workspace, bool grad, bool accumulate,
+                                bool use_split_accumulator, cudaStream_t stream_main) {
+#ifdef USE_HIPKITTENS_GEMM
+  if (kittens_fused_ag_gemm_supported(cuda::sm_arch())) {
+    const bool launched = hk_bulk_ag_gemm(A, transa, B, transb, D, bias, pre_gelu_out, workspace,
+                                          accumulate, _ubuf, _ubufs[0], _ub_comm, _ub_reg, _tp_id,
+                                          _tp_size, _ag_signal_base + _tp_size, stream_main);
+    NVTE_CHECK(launched, "fused bulk AG failed to launch");
+    _ag_signal_base += _tp_size;
+    return;
+  }
+#endif
+  NVTE_ERROR("fused bulk AG was selected but is not built into this library");
+}
+
 // TODO: Introduce HIPGraphs for dependency management.
 void CommOverlapP2PBase::rocm_split_overlap_ag(const TensorWrapper &A, bool transa, const TensorWrapper &B,
                                 bool transb, TensorWrapper &D, TensorWrapper &bias,