Replaced cublasLt calls with rocblas gemm + 6 manually crafted epilogues

[wenchenvincent]

No description provided.

[HaiShaw]

Have some concern regarding the order of ops in backward pass, in case both bias and gelu are invoked. Minor but careful concerns are over bias dimensional sizes.
Very good in general!

[HaiShaw]

This assumes the row major storage, which is different than the default of remaining libraries.
I think if we change it to thread_data = (float)in[row_idx + m * col_idx] then it is easier to understand and we can unbend the confusing comment and code like belows:

   // The bias vector length is m. So it will be reduced along axis 0 in row major
   ... ...
  NVTE_CHECK_CUDA( hipMalloc(&bias_tmp, sizeof(float)*m) );   // L398
  ... ...
  NVTE_CHECK_CUDA( hipMalloc(&bias_tmp, sizeof(float)*n) );   // L450

It is confusing is because bias vector should always be length of n (from a same context of a m*n matrix), if we treat the storage formation consistently cross APIs.

[HaiShaw]

Btw, we may want to unstrict us from using (float) only, onwards.
This BlackReduce seems to be templated. But it is fine for now.

[wenchenvincent]

Here, the data layout is indeed row major as it is C/C++ convention. Only rocBlas assumes column major layout. Outside rocBlas, it is always row major.

Notice that the perception that the bias vector should always be length of n assumes that n is the number of output features. However, when we call GEMM, m, k, n have been changed and it is not always the case that n is the number of the output features.

[wenchenvincent]

It seems that in many kernels of NVTE, the computation is done in fp32, while the input/output could be of lower precision. Especially for reduction operations, fp32 seems to be preferred.

[HaiShaw]

Regarding the row or column major, we are not dealing with native HOST C/C++ arrays, so what C/C++ compiler uses for that is irrelevant. As far as I know, beyond rocBLAS, CK, cuBLAS, BLAS, most math libraries took column majors as default.
I'm not pointing logic errors, but mostly for code clarity here. As it stands, I see inconsistent hipMalloc - hard to follow or looks unclear to me. I'm proposing a consistent view only, hoping to make logic more clear though.

[HaiShaw]

However, when we call GEMM, m, k, n have been changed and it is not always the case that n is the number of the output features.
It would be helpful with some comment s to tell when m & n may have changed and how.
It seems hard to follow inconsistent code like below:

  NVTE_CHECK_CUDA( hipMalloc(&bias_tmp, sizeof(float)*m) );   // L398
  ... ...
  NVTE_CHECK_CUDA( hipMalloc(&bias_tmp, sizeof(float)*n) );   // L450

[wenchenvincent]

I am not sure if most math libraries take column major as default. rocBlas assumes the data is column major. But in frameworks, the data is arranged in row major.

Also, the main reason why the bias vector length is sometimes m and sometimes n, is that the assumption that n is the number of output features no longer holds true in some GEMM configs.

[wenchenvincent]

Here is the explanation regarding why in CUBLASLT_EPILOGUE_BIAS, the bias length is GEMM_m while in CUBLASLT_EPILOGUE_BGRADB, the bias length is GEMM_n.

For a linear layer in PyTorch, we have the following equation for forward pass:

Y = X * W^T + b. X is of shape (m, k), W is of shape (n, k) and Y is of shape (m,n).

m - batch size
k - input feature dimension
n - output feature dimension.

Notice that m, k, n here are not necessarily the m, k, n used in GEMM calls. So we will denote the m, k, n used in GEMM as GEMM_m, GEMM_k, GEMM_n later.

For the forward pass, we want to compute Y = X * W^T + b, and the data in X and W are row major. But cuBlas/cuBlasLt/rocBlas assumes the data to be column major. To avoid the hassle of converting the data layout from row major to column major, what we can do is to call GEMM(W, X) with transa=T.

Here W is of shape (n,k) in row major and of shape (k, n) in column major.
X is of shape (m, k) in row major and of shape (k, m) in column major.
The result is of shape (m, n) in row major and of shape (n, m) in column major.

In the GEMM call for the forward pass, matrix A is W (with transa=T), matrix B is X, so

GEMM_m = n
GEMM_k = k
GEMM_n = m

We can see that the bias vector length n is actually GEMM_m. That's why in CUBLASLT_EPILOGUE_BIAS, we say the bias vector length is "m". GEMM call itself is not aware of the batch size, input feature dimension, or output feature dimension in the NN layer. All it knows is GEMM_m, GEMM_k, and GEMM_n.

For the backward pass, we have:

dX = dY * W
dW = dY^T * X
db = sum(dY, axis=0)

For CUBLASLT_EPILOGUE_BIAS, they compute db along with dW.

Again, notice that only when we use cuBlas/cuBlasLt/rocBlas will data be treated as column-major layout. In other contexts, they are row major. So to compute dW without actually converting the input matrices from row-major to column-major layout, we will call rocblas_gemm_ex(X, dY), with transb=T.

Here X is of shape (m, k) in row major and of shape (k, m) in column major.
dY is of shape (m, n) in row major and of shape (n, m) in column major.
The result is of shape (k, n) in column major and of shape (n, k) in row major, which is exactly dW

So in this gemm call, the matrix B is dY, and transb=T, and the matrix D is dW. db is computed by reducing dY.
We can denote the m,n,k in the GEMM for computing dW as gemm_m, gemm_n, and gemm_k. We have the following:

gemm_m = k
gemm_k = m
gemm_n = n

So here the bias vector length n is gemm_n. That's why in CUBLASLT_EPILOGUE_BIAS, the bias length is "n".

[HaiShaw]

This explain is perfectly correct!
I think regarding the code itself, we can add a comment to each case (within our cublas_gemm wrapper) with something like:

// gemm_ex(A=W, B=X), determined by user code in module.py
// gemm_ex(A=X, B=dy), determined by user code in module.py

to make it easier to read.
We may also define (to each different case)

math_m = m 
math_k = k
math_n = n

or

math_m = n
math_k = k
math_n = m

or 

math_m = k
math_k = m
math_n = n

assume m/k/n the gemm_m/gemm_k/gemm_n within the wrapper function.

This way (together with previous comment added) we will have an uniform and easier comprehensive code at this layer, for something like:

NVTE_CHECK_CUDA( hipMalloc(&bias_tmp, sizeof(float)*math_n) );

[HaiShaw]

LGTM

[HaiShaw]

I assume previous 'n' should be 'm'. No question to the modified