Repository: NVIDIA/TransformerEngine

Issue:
CUBLAS_STATUS_EXECUTION_FAILED when calling cublasGemmStridedBatchedExFix
I have an LLM completely rewritten to use TransformerEngine's layers. I'm getting this error on forward (doing inference)

```python
  File "/usr/src/app/lit-gpt/lit_gpt/model_all_replaced.py", line 231, in forward
    x = x + h + self.norm_2_mlp(x)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1502, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1511, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/transformer_engine/pytorch/module/layernorm_mlp.py", line 1222, in forward
    out = fwd_fn(*args)
  File "/usr/local/lib/python3.8/dist-packages/transformer_engine/pytorch/module/layernorm_mlp.py", line 343, in forward
    gelu_out = bias_gelu_fused(fc1_out, fc1_bias)
  File "/usr/local/lib/python3.8/dist-packages/transformer_engine/pytorch/jit.py", line 94, in bias_gelu_fused
    return bias_gelu_fused_(inp, bias)
  File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/eval_frame.py", line 295, in _fn
    return fn(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/eval_frame.py", line 448, in catch_errors
    return callback(frame, cache_size, hooks, frame_state)
  File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/convert_frame.py", line 526, in _convert_frame
    result = inner_convert(frame, cache_size, hooks, frame_state)
  File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/convert_frame.py", line 134, in _fn
    torch.cuda.set_rng_state(cuda_rng_state)
  File "/usr/local/lib/python3.8/dist-packages/torch/cuda/random.py", line 65, in set_rng_state
    _lazy_call(cb)
  File "/usr/local/lib/python3.8/dist-packages/torch/cuda/__init__.py", line 202, in _lazy_call
    callable()
  File "/usr/local/lib/python3.8/dist-packages/torch/cuda/random.py", line 63, in cb
    default_generator.set_state(new_state_copy)
RuntimeError: CUDA error: misaligned address
```

With `CUDA_LAUNCH_BLOCKING=1`

```python
Traceback (most recent call last):
  File "repro_bug.py", line 13, in <module>
    y = m(x, padding_multiple=8, max_seq_length=65, input_pos=input_pos)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1502, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1511, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/src/app/lit-gpt/lit_gpt/model_all_replaced.py", line 120, in forward
    x, self.kv_caches[i] = block(x, (cos, sin), max_seq_length, mask, input_pos, self.kv_caches[i], padding)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1502, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1511, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/src/app/lit-gpt/lit_gpt/model_all_replaced.py", line 224, in forward
    y = self.attn(q, k, v, attention_mask=(~mask if mask is not None else None))
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1502, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1511, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/transformer_engine/pytorch/attention.py", line 946, in forward
    return self.unfused_attention(query_layer, key_layer, value_layer, attention_mask)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1502, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1511, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/transformer_engine/pytorch/attention.py", line 242, in forward
    context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1))
RuntimeError: CUDA error: CUBLAS_STATUS_EXECUTION_FAILED when calling `cublasGemmStridedBatchedExFix(handle, opa, opb, (int)m, (int)n, (int)k, (void*)&falpha, a, CUDA_R_16BF, (int)lda, stridea, b, CUDA_R_16BF, (int)ldb, strideb, (void*)&fbeta, c, CUDA_R_16BF, (int)ldc, stridec, (int)num_batches, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)`
```

The model is https://github.com/Lightning-AI/lit-gpt/blob/carmocca/h100/lit_gpt/model_all_replaced.py

Here's how you can reproduce it

```python
git clone https://github.com/Lightning-AI/lit-gpt && cd lit-gpt && git checkout carmocca/h100
```

run:

```python
import lit_gpt.model_transformer_engine as model
import torch

model.ATTN_MASK_TYPE = "padding"

torch.set_default_dtype(torch.bfloat16)

with torch.device("cuda"):
    m = model.GPT.from_name("pythia-70m")
    x = torch.randint(0, 5, (1, 5), dtype=torch.int)
    input_pos = torch.tensor([0, 1, 2, 3, 4], dtype=torch.long)

# changing these will circumvent the issue
y = m(x, padding_multiple=8, max_seq_length=65, input_pos=input_pos)
```

Note that this ran on an A100 machine with CUDA 12.0. I'm having the problem on an H100 machine:

```bash
root@h100:/usr/src/app# nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2023 NVIDIA Corporation
Built on Fri_Jan__6_16:45:21_PST_2023
Cuda compilation tools, release 12.0, V12.0.140
Build cuda_12.0.r12.0/compiler.32267302_0
root@h100:/usr/src/app# nvidia-smi
Fri Jun 23 16:29:00 2023       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.105.17   Driver Version: 525.105.17   CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA H100 80G...  On   | 00000000:1C:00.0 Off |                    0 |
| N/A   34C    P0   110W / 700W |  26773MiB / 81559MiB |      6%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
+-----------------------------------------------------------------------------+
```

Thank you for your help!


Discussion:
@carmocca Thanks for creating the issue! Could you post the output of `python -m torch.utils.collect_env`, please? 
Unfortunately, I dont have access to the H100 machine anymore, but I used:

ghcr.io/coreweave/ml-containers/torch:ceeb8c2-nccl-cuda12.0.1-nccl2.18.1-1-torch2.0.1-vision0.15.2-audio2.0.2

For the image base

And installed PyTorch nightly and TransformerEngine main
Do things work if you change the sequence length from 65 to a multiple of 16? cuBLAS requires FP8 matrices to be aligned to 16B, so I wonder if that is causing problems.
@carmocca I would need more information on how to reproduce the environment.
Inside my setup I'm running:
```
git clone https://github.com/Lightning-AI/lit-gpt && cd lit-gpt && git checkout carmocca/h100
python setup.py install
pip install -r ./requirements.txt
# create a main.py script with your Python code
python main.py
```
which results in:
```
Traceback (most recent call last):
  File "/workspace/src/lit-gpt/main.py", line 1, in <module>
    import lit_gpt.model_all_replaced as model
ModuleNotFoundError: No module named 'lit_gpt.model_all_replaced'
```
so I assume I'm not using the right install command for your branch. 

EDIT: I see you have linked the missing file here: https://github.com/Lightning-AI/lit-gpt/blob/carmocca/h100/lit_gpt/model_all_replaced.py which returns a `404`. 
I'm sorry @ptrblck, I renamed the file for my PR but didn't think of updating this repro. The link is now https://github.com/Lightning-AI/lit-gpt/blob/carmocca/h100/lit_gpt/model_transformer_engine.py. I updated the top post.

For context, this is my branch for https://github.com/Lightning-AI/lit-gpt/pull/123
> Do things work if you change the sequence length from 65 to a multiple of 16? cuBLAS requires FP8 matrices to be aligned to 16B, so I wonder if that is causing problems.

The implementation pads the input to fulfill that requirement, see `padding_multiple=8` argument in the repro script. Interestingly, the repro doesn't fail for `max_seq_length` values of less than `65`
@timmoon10 the error comes from `torch.bmm` inside attention so fp8 is not involved here.
The repro script is not very minimal, but the time I had available on an H100 machine was limited.
For this reason I posted the issue report as is, considering that maybe a maintainer would isolate the issue further.

I do know that the issue didnt happen on an A100. The same (numerically) model without the TransformerEngine layers also doesnt have the issue.

If you still consider this a 3rd party issue, we can close it and I'll reopen it again in the appropiate repo when I get access to the hardware again.


@carmocca Thanks for the great repro!
I've isolated this issue to the `FusedScaleMaskSoftmax` kernel in TE. Basically it appears that this kernel doesn't handle the exact shape provided correctly, incurs an illegal memory access (in the form of the warp misaligned address), and then `cuBLAS` is surfacing the failure as it is attempting to launch the next kernel in a corrupted CUDA context. I'm fairly confident that cuBLAS is not to blame here as running the same cuBLAS workload in isolation (down to the tensor strides) doesn't surface any failures.

I've also verified that you can temporarily work around the issue by bypassing the TE softmax kernel for this case by e.g., editing the condition in https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/softmax.py#L269 to avoid using TE's softmax kernel for this case.

Resolving patch:
diff --git a/transformer_engine/pytorch/softmax.py b/transformer_engine/pytorch/softmax.py
--- a/transformer_engine/pytorch/softmax.py
+++ b/transformer_engine/pytorch/softmax.py
@@ -270,10 +270,11 @@ def is_kernel_available(self, b: int, np: int, sq: int, sk: int) -> bool:
         """Check FusedScaleMaskSoftmax kernel availability based on size"""
         attn_batches = b * np
 
-        if (
+        if ( # pylint: disable=too-many-boolean-expressions
             self.scaled_masked_softmax_fusion  # user want to fuse
             and self.input_in_float16  # input must be fp16
             and 16 < sk <= 4096  # sk must be 16 ~ 2048
+            and sk % 8 == 0  # sk must be divisor of 8
             and sq % 4 == 0  # sq must be divisor of 4
             and attn_batches % 4 == 0  # np * b must be divisor of 4
         ):


Test patch: