InternVL2-8B and InternVL2-26B CUDA error: an illegal memory access was encountered

[HubHop]

Hi authors,

Was trying to run InternVL-8B and InternVL-26B on 4 GPUs, but I got this,

  File ".cache/huggingface/modules/transformers_modules/main/modeling_internlm2.py", line 656, in forward
    hidden_states, self_attn_weights, present_key_value = self.attention(
  File "...python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File ".cache/huggingface/modules/transformers_modules/main/modeling_internlm2.py", line 498, in forward
    attn_output = self._flash_attention_forward(
  File ".cache/huggingface/modules/transformers_modules/main/modeling_internlm2.py", line 535, in _flash_attention_forward
    query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._unpad_input(
  File ".cache/huggingface/modules/transformers_modules/main/modeling_internlm2.py", line 564, in _unpad_input
    indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
  File ".cache/huggingface/modules/transformers_modules/main/modeling_internlm2.py", line 85, in _get_unpad_data
    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
RuntimeError: CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Code is exactly the same as in huggingface demo.

My python env:

• flash-attn: 2.5.6
• transformers: 4.37.2
• accelerate: 0.32.1
• torch: 2.0.1+cu117

I guess the error is due to flash attention, but not 100% sure. Any help?

Thanks in advance!

[czczup]

Could you try this split_model function:

def split_model(model_name):
    device_map = {}
    world_size = torch.cuda.device_count()
    num_layers = {'InternVL2-8B': 32, 'InternVL2-26B': 48,
                  'InternVL2-40B': 60, 'InternVL2-Llama3-76B': 80}[model_name]
    # Since the first GPU will be used for ViT, treat it as half a GPU.
    num_layers_per_gpu = math.ceil(num_layers / (world_size - 0.5))
    num_layers_per_gpu = [num_layers_per_gpu] * world_size
    num_layers_per_gpu[0] = math.ceil(num_layers_per_gpu[0] * 0.5)
    layer_cnt = 0
    for i, num_layer in enumerate(num_layers_per_gpu):
        for j in range(num_layer):
            device_map[f'language_model.model.layers.{layer_cnt}'] = i
            layer_cnt += 1
    device_map['vision_model'] = 0
    device_map['mlp1'] = 0
    device_map['language_model.model.tok_embeddings'] = 0
    device_map['language_model.model.embed_tokens'] = 0
    device_map['language_model.output'] = 0
    device_map['language_model.model.norm'] = 0
    device_map['language_model.lm_head'] = 0
    device_map[f'language_model.model.layers.{num_layers - 1}'] = 0

    return device_map


path = 'OpenGVLab/InternVL2-26B'
device_map = split_model('InternVL2-26B')
print(device_map)
model = AutoModel.from_pretrained(
    path,
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
    trust_remote_code=True,
    device_map=device_map).eval()

[HubHop]

Thanks for the help @czczup, I just tried this new code and unfortunately it still doesn't work. My setting is based on 8 GPUs and I can run other scales (e.g., 1B, 4B, 40B, 76B) successfully.

[czczup]