Wrong metrics for FLOP utilization and TensorCore usage when using bfloat16

[emergenz]

The following program is a matrix multiplication of two matrices in BF16 on an H100. We achieve a rough throughput of 600 TFLOP/s (which is impossible without using TensorCores). BF16 matrix multiplication is TensorCore-eligible on the H100. The GPU kernel stats page of the JAX profiler incorrectly states that the op is not TensorCore-eligible and that no TensorCores are used. The framework op stats page also incorrectly identifies the op as not being TensorCore-eligible.

However, again, such high throughput is impossible without using TensorCores (I stated the throughput manually measured by myself, which corresponds to the one listed in the graph viewer).

On a slightly different note, for the FLOPS utilization it is unclear to me whether the peak FLOP/s used for the calculation are correct (which should be different for BF16 and TF32).

[emergenz]

I just realized that I didn't append the actual program to this issue. My bad. It was exactly the program in #1704 but with the inputs cast to BF16.

import datetime
import jax.numpy as jnp
import jax

MATRIX_DIM = 32768
STEPS = 10

A = jnp.ones((MATRIX_DIM, MATRIX_DIM), dtype=jnp.bfloat16)
B = jnp.ones((MATRIX_DIM, MATRIX_DIM), dtype=jnp.bfloat16)

num_bytes = A.size * 2
total_num_bytes_crossing_to_hbm = num_bytes * 3

total_num_flops = 2 * MATRIX_DIM * MATRIX_DIM**2

def matmul(A, B):
  return A @ B

matmul(A, B) # warmup

with jax.profiler.trace("tensorboard"):
  start_time = datetime.datetime.now()
  for i in range(STEPS):
    C = A @ B
    C.block_until_ready()
  end_time = datetime.datetime.now()

average_time_per_step = (end_time - start_time).total_seconds() / STEPS

print(f"{average_time_per_step}, teraflops per second: {total_num_flops / average_time_per_step / 1e12}, gigabytes per second: {total_num_bytes_crossing_to_hbm / average_time_per_step / 1e9}")

[Blair-Johnson]

It looks like the profiler uses string matching to identify some common patterns of tensorcore-utilizing kernels.

xprof/xprof/utils/kernel_stats_utils.cc

Lines 43 to 65 in e7edebd

	// A list of patterns to help determine if a kernel uses Tensor Core.
	// A kernel uses Tensor Core if its kernel name contains any of these patterns.
	// Some examples of kernel names: volta_h884gemm, turing_fp16_s1688cudnn_fp16
	constexpr absl::string_view kTensorCoreKernelNamePatterns[] = {
	"16816",
	"c1688",
	"conv1x1",
	"conv2d_c1_k1",
	"dgrad_1x1_stride_2x2",
	"direct_group",
	"first_layer_wgrad_kernel",
	"h1688",
	"h884",
	"hmma",
	"i16832",
	"i8816",
	"s884",
	"s1688",
	"xmma_gemm",
	"xmma_implicit_gemm",
	"xmma_sparse_conv",
	"xmma_sparse_gemm",
	"xmma_warp_specialized_implicit_gemm"};

As a user, this also caused me significant confusion. So people don't go down rabbit-holes trying to figure out why their code isn't hitting the tensor-cores, it might be good to add an asterisk in the UI or otherwise flag that this isn't "ground truth".

[Matt-Hurd]

@cliveverghese can you take a look at this and see if we can reasonably explain this to the user?