demo.py RuntimeError: CUDA error: an illegal memory access was encountered

[Cikguseven]

I set up the repo using Option 1 and downloaded the model using download_blt_weights.py. I then ran demo.py with the following prompt from the Physical Interaction: Question Answering (PIQA) dataset

python demo.py "Question: How do you properly prepare a steak.
Answer: Take the steak out of warm storage and let come to room temperature, generously add salt and pepper to both sides and let sit for 10 minutes."

and got the error:

╭────────────────────────────────────────────────────────────────────────────── locals ──────────────────────────────────────────────────────────────────────────────╮ │
[rank0]: │ │ attn_bias_type = 'local_block_causal'                                                                                                                              │ │
[rank0]: │ │      attn_impl = 'xformers'                                                                                                                                        │ │
[rank0]: │ │         eos_id = 2                                                                                                                                                 │ │
[rank0]: │ │         seqlen = 260                                                                                                                                               │ │
[rank0]: │ │ sliding_window = 512                                                                                                                                               │ │
[rank0]: │ │         tokens = <repr-error 'CUDA error: an illegal memory access was encountered\nCUDA kernel errors might be asynchronously reported at some other API call, so │ │
[rank0]: │ │                  the stacktrace below might be incorrect.\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1\nCompile with `TORCH_USE_CUDA_DSA` to enable      │ │
[rank0]: │ │                  device-side assertions.\n'>                                                                                                                       │ │
[rank0]: │ ╰────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ │
[rank0]: │                                                                                                                                                                        │
[rank0]: │ /home/kieron/blt/bytelatent/model/utils.py:116 in tokens_to_seqlen                                                                                                     │
[rank0]: │                                                                                                                                                                        │
[rank0]: │   113 │                                                                                                                                                                │
[rank0]: │   114 │   # 0 0 0 1 0 0 0 1 0 0 X                                                                                                                                      │
[rank0]: │   115 │   # 0 1 0 0 0 1 0 0 0 0 X                                                                                                                                      │
[rank0]: │ ❱ 116 │   row, col = torch.where(mask)                                                                                                                                 │
[rank0]: │   117 │                                                                                                                                                                │
[rank0]: │   118 │   # row = 0, 0, 0, 1, 1, 1                                                                                                                                     │
[rank0]: │   119 │   # col = 3, 7, 10, 1, 5, 10                                                                                                                                   │
[rank0]: │                                                                                                                                                                        │
[rank0]: │ ╭────────────────────────────────────────────────────────────────────────────── locals ──────────────────────────────────────────────────────────────────────────────╮ │
[rank0]: │ │  batch = <repr-error 'CUDA error: an illegal memory access was encountered\nCUDA kernel errors might be asynchronously reported at some other API call, so the     │ │
[rank0]: │ │          stacktrace below might be incorrect.\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1\nCompile with `TORCH_USE_CUDA_DSA` to enable device-side      │ │
[rank0]: │ │          assertions.\n'>                                                                                                                                           │ │
[rank0]: │ │ eos_id = 2                                                                                                                                                         │ │
[rank0]: │ │   mask = <repr-error 'CUDA error: an illegal memory access was encountered\nCUDA kernel errors might be asynchronously reported at some other API call, so the     │ │
[rank0]: │ │          stacktrace below might be incorrect.\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1\nCompile with `TORCH_USE_CUDA_DSA` to enable device-side      │ │
[rank0]: │ │          assertions.\n'>                                                                                                                                           │ │
[rank0]: │ ╰────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ │
[rank0]: ╰────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
[rank0]: RuntimeError: CUDA error: an illegal memory access was encountered
[rank0]: CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
[rank0]: For debugging consider passing CUDA_LAUNCH_BLOCKING=1
[rank0]: Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Aborted (core dumped)

[kofinandi]

Same issue on RTX 3090 or A6000 single-GPU setup.

Changing to _compile=False, in blt.py in the following lines fixed it, and the demo runs as expected.

def patch_mask(b, h, q_idx, kv_idx):
    return cross_mask[b, q_idx, kv_idx]

block_mask = create_block_mask(
    patch_mask,
    B=bs,
    H=None,
    Q_LEN=q_len,
    KV_LEN=kv_len,
    _compile=True, # change this to False to fix
)

Or switching to an A100 GPU also makes the original _compile=True, work.

Seems like a PyTorch issue, reported it here: pytorch/pytorch#166070

[Cikguseven]

@kofinandi _compile=True does not seem to work on both 40GB and 80GB A100 for me. I'm also using torch 2.6.0.dev20241112+cu121.