Fp8 gemm enabled by wenchenvincent · Pull Request #7 · ROCm/TransformerEngine

wenchenvincent · 2023-04-15T04:43:59Z

Currently supported fp8 input and fp32 output.
fp8 output not supported yet.

…to NANOO mode.

HaiShaw

LGTM with minor concerns.
Importantly - we should give f8 a refresh to reflect ongoing support of NANOO and OCP only, so suggest to remove IEEE and related dead code, insert OCP instead.

HaiShaw · 2023-08-22T19:34:42Z

+		    2. input: fp32, output: fp32.\n 3. input: fp8, output: fp32");
+
+
+    //If D is not fp32, then we need a temp buffer for GEMM result before applying epilogues. Otherwise, we can apply epilogues in-place.


May want to consider to refactor this long logic expression (to a branch of subconditions) for better avg. performance

Not sure what you mean by a branch of subconditions.

HaiShaw · 2023-08-22T19:42:26Z

+    //If D is not fp32, then we need a temp buffer for GEMM result before applying epilogues. Otherwise, we can apply epilogues in-place.
    void* D_temp;
-    if ((bias || gelu) && (A_type==rocblas_datatype_f16_r && B_type==rocblas_datatype_f16_r && D_type==rocblas_datatype_f16_r)) {
+    if ((bias || gelu) && (D_type==rocblas_datatype_f16_r || D_type==rocblas_datatype_f8_r || D_type==rocblas_datatype_bf8_r)) {


lines above NVTE_CHECK made sure D_type is not f8/bf8 (for the moment), so here the later 2 conditions seem to be stale, but it is minor.

Yeah, I add these two conditions in advance as there would be support for fp8 output in the future. Right now, it is limited to fp32/fp16 output.

HaiShaw · 2023-08-22T19:48:32Z

    // TODO: Can we search for rocblas_gemm_algo??
-    NVTE_CHECK_CUBLAS(rocblas_gemm_ex(handle, transa, transb, m, n, k, &one,
+    if (use_fp8) {
+      rocblas_computetype computeType = rocblas_compute_type_f32;


Is it rocblas_compute_type_f32_r?

No, it is rocblas_compute_type_f32

HaiShaw · 2023-08-22T20:00:41Z

        }
    }
-    if ((bias || gelu) && (A_type==rocblas_datatype_f16_r && B_type==rocblas_datatype_f16_r && D_type==rocblas_datatype_f16_r)) {
+    if ((bias || gelu) && (D_type==rocblas_datatype_f16_r || D_type==rocblas_datatype_f8_r || D_type==rocblas_datatype_bf8_r)) {


same, D_type should not be f8/bf8 from earlier check, minor.

HaiShaw · 2023-08-22T20:10:22Z

    """

-    E4M3 = _FormatHelper(max_fwd=448, max_bwd=448)
+    E4M3 = _FormatHelper(max_fwd=240, max_bwd=240)


We should make our cap at 224 vs. 240
Also, this is based on assumption that IEEE mode is phasing out.
We should refresh on hip_f8 headers to reflect IEEE's phase out, and to support NANOO and OCP only.

We haven't made the decision on clipping at 224 yet, right (though it was discussed multiple times)? Once we make our final decisions, we can change it.

We will remove IEEE mode. We don't really have support for OCP yet. Will add it once it is time.

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Applied review comments 1, 2, 3, 4, 6, 7, 8 from PR #667 in one pass. Comments 5 and 9 will be answered on-thread (see below). Applied ======= #7 cpp_extensions/gemm.py: gate NVTE_USE_GEMM_TRITON behind IS_HIP_EXTENSION. Our Triton kernels use gfx942/gfx950-specific MFMA instructions and autotune configs; refuse to enable on non-HIP builds. #1 + #4 + #8: deduplicate dtype utilities against triton_kernels/common.py, which is the authoritative source shared across all Triton kernel backends here. - gemm_common.py: drop our copies of _get_fp8_dtypes(), torch_to_te_dtype(), te_to_torch_dtype(). Keep is_fp8_dtype() (operates on tex.DType; common.py's is_fp8_torch_dtype operates on torch dtypes -- different signatures), reinterpret_as_fp8_tensor, getGemmOutputShape, product. - gemm_wrapper.py: import torch_dtype_to_te_dtype / te_dtype_to_torch_dtype from ..common; use those instead of the removed local names. - __init__.py: drop the redundant re-exports (torch_to_te_dtype, te_to_torch_dtype, _get_fp8_dtypes). Note in the file directing callers to triton_kernels.common. - test_gemm_kernel.py: import get_torch_e4m3_type, get_torch_e5m2_type, torch_dtype_to_te_dtype from triton_kernels.common; define a local _get_fp8_dtypes() shim over them to keep the test's parametrization tables unchanged. - test_gemm_mxfp8.py: previously used _get_e4m3_dtype() -- dropped along with the test that used it (see below). #2 + #3: add module docstrings to test_gemm.py, test_gemm_kernel.py, test_gemm_mxfp8.py making the scope of each file explicit and pointing at the sibling files. Prompted by wangye805's confusion between test_gemm.py and test_gemm_kernel.py, and by the question of whether MXFP8 is only tested in test_gemm_mxfp8.py. #6 test_numerics.py: rewrite the comment above and the reason= text for _skip_grouped_under_gemm_triton to make it explicit that the skip is a *backend mismatch* between the two sides of the comparison, not Triton non-determinism. - sequential side -> our Triton (via NVTE_USE_GEMM_TRITON=1) - grouped side -> hipBLASLt / CUTLASS / AITER-Triton grouped (controlled by the SEPARATE NVTE_USE_GROUPED_GEMM_TRITON env var) With just NVTE_USE_GEMM_TRITON=1 set, the two sides diverge in fp32 rounding because they run different kernels. Not our Triton being non-deterministic. Also removed test_mxfp8_kernel_with_simulated_data (an early kernel- bring-up smoke test that only asserted non-zero output; MXFP8 correctness is fully covered by test_gemm.py::test_triton_vs_*_mxfp8 with real MXFP8Tensor and both PyTorch/C++ references). To be answered on-thread ======================== #5 (conftest.py hook vs explicit skip lists) -- keep the hook: it fires only on three specific ValueError substrings from our own gate code, is self-retiring when the gates are relaxed, and explicit skips would require ~1300 marks. Detail in the PR reply. #9 (Float8TensorWrapper / MXFP8TensorWrapper vs the TE tensor types) -- defer; needs a side-by-side to see how much of the wrapper API is genuinely required for the Triton-kernel call surface (fields accessed, dimension reordering, missing-transpose fallbacks) vs. what could route through the TE tensor types directly. Verified on gfx950 / PyTorch 2.10 -- all three direct suites clean: test_gemm.py 212 pass / 72 skip test_gemm_kernel.py 450 pass / 606 skip test_gemm_mxfp8.py 2 pass (was 3; smoke test removed)

wenchenvincent added 2 commits April 5, 2023 04:31

Initial integration of rocBlas F8 GEMM without epilogues. Changed F8 …

32045db

…to NANOO mode.

Added support for bias data type different from matrix B or D.

8a4d985

wenchenvincent requested a review from HaiShaw April 15, 2023 04:43

HaiShaw approved these changes Aug 22, 2023

View reviewed changes

wenchenvincent merged commit ba0be60 into dev Aug 22, 2023

wenchenvincent pushed a commit that referenced this pull request Sep 17, 2024

rm deprecated pytest-runner requirement from setup (#7)

2db8376

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Fp8 gemm enabled#7

Fp8 gemm enabled#7
wenchenvincent merged 2 commits into
devfrom
fp8_gemm

wenchenvincent commented Apr 15, 2023

Uh oh!

HaiShaw left a comment

Uh oh!

HaiShaw Aug 22, 2023

Uh oh!

wenchenvincent Aug 22, 2023

Uh oh!

HaiShaw Aug 22, 2023

Uh oh!

wenchenvincent Aug 22, 2023

Uh oh!

HaiShaw Aug 22, 2023

Uh oh!

wenchenvincent Aug 22, 2023

Uh oh!

HaiShaw Aug 22, 2023

Uh oh!

HaiShaw Aug 22, 2023

Uh oh!

wenchenvincent Aug 22, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		2. input: fp32, output: fp32.\n 3. input: fp8, output: fp32");


		//If D is not fp32, then we need a temp buffer for GEMM result before applying epilogues. Otherwise, we can apply epilogues in-place.

Uh oh!

Conversation

wenchenvincent commented Apr 15, 2023

Uh oh!

HaiShaw left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants