RuntimeError: value cannot be converted to type uint8_t without overflow with masked_fill in GPT-Neo causal mask creation on DirectML. torch.where works.

[Freeaker1]

When running inference with the EleutherAI/gpt-neo-1.3B model from Hugging Face Transformers on the DirectML backend, a RuntimeError: value cannot be converted to type uint8_t without overflow occurs.
The error traces back to the masked_fill operation within the _prepare_4d_causal_attention_mask_with_cache_position method in transformers.models.gpt_neo.modeling_gpt_neo.py.
The model and script run correctly on the CPU backend.

This issue was debugged with the assistance of an AI. cc Matt Todd (@mtodd) (Matt Todd), Ian Baird (@ijbaird) (Ian Baird) (if relevant).

Pasos para Reproducir:

1. Set up a Python environment with the specified versions (see "Entorno" section below).
2. Run the provided inference script (see "Script de Inferencia" section below) with the model EleutherAI/gpt-neo-1.3B targeting the DirectML device.

Comportamiento Actual (Error con masked_fill):

The script fails with the following error and traceback:

!!!! DEBUG_MASKED_FILL: ERROR during masked_fill operation: value cannot be converted to type uint8_t without overflow !!!!
Traceback (most recent call last):
  File "D:\LLaMA3\llama_env\Lib\site-packages\transformers\models\gpt_neo\modeling_gpt_neo.py", line 956, in _prepare_4d_causal_attention_mask_with_cache_position
    causal_mask_clone[:, :, :, :mask_length] = target_for_masked_fill.masked_fill(
                                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: value cannot be converted to type uint8_t without overflow
Error durante la inferencia: value cannot be converted to type uint8_t without overflow
Traceback (most recent call last):
  File "D:\LLaMA3\inferencia.py", line 92, in <module>
    output = model.generate(
             ^^^^^^^^^^^^^^^
  File "D:\LLaMA3\llama_env\Lib\site-packages\torch\utils\_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "D:\LLaMA3\llama_env\Lib\site-packages\transformers\generation\utils.py", line 2465, in generate
    result = self._sample(
             ^^^^^^^^^^^^^
  File "D:\LLaMA3\llama_env\Lib\site-packages\transformers\generation\utils.py", line 3431, in _sample
    outputs = self(**model_inputs, return_dict=True)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\LLaMA3\llama_env\Lib\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\LLaMA3\llama_env\Lib\site-packages\torch\nn\modules\module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\LLaMA3\llama_env\Lib\site-packages\transformers\models\gpt_neo\modeling_gpt_neo.py", line 1027, in forward
    transformer_outputs = self.transformer(
                          ^^^^^^^^^^^^^^^^^
  File "D:\LLaMA3\llama_env\Lib\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\LLaMA3\llama_env\Lib\site-packages\torch\nn\modules\module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\LLaMA3\llama_env\Lib\site-packages\transformers\models\gpt_neo\modeling_gpt_neo.py", line 715, in forward
    causal_mask = self._update_causal_mask(
                  ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\LLaMA3\llama_env\Lib\site-packages\transformers\models\gpt_neo\modeling_gpt_neo.py", line 841, in _update_causal_mask
    causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\LLaMA3\llama_env\Lib\site-packages\transformers\models\gpt_neo\modeling_gpt_neo.py", line 965, in _prepare_4d_causal_attention_mask_with_cache_position 
    raise e_masked_fill_debug
  File "D:\LLaMA3\llama_env\Lib\site-packages\transformers\models\gpt_neo\modeling_gpt_neo.py", line 956, in _prepare_4d_causal_attention_mask_with_cache_position
    causal_mask_clone[:, :, :, :mask_length] = target_for_masked_fill.masked_fill(
                                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: value cannot be converted to type uint8_t without overflow

Workaround Encontrado:

The issue is resolved by modifying transformers/models/gpt_neo/modeling_gpt_neo.py (specifically the _prepare_4d_causal_attention_mask_with_cache_position method). Replacing the problematic masked_fill line with an equivalent torch.where operation allows the inference to complete successfully on DirectML.

The original logic within _prepare_4d_causal_attention_mask_with_cache_position involving masked_fill:

# target_for_masked_fill = causal_mask_clone[:, :, :, :mask_length]
# padding_mask = ((causal_mask_clone[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(causal_mask_clone.device)) == 0)
# min_dtype_val = torch.finfo(dtype).min

# Problematic line:
# causal_mask_clone[:, :, :, :mask_length] = target_for_masked_fill.masked_fill(
#     padding_mask, min_dtype_val
# )

Was replaced with torch.where:

# target_for_op = causal_mask_clone[:, :, :, :mask_length]
# padding_mask_bool = ((causal_mask_clone[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(causal_mask_clone.device)) == 0) # Simplified for brevity
# min_dtype_val = torch.finfo(dtype).min

filled_slice = torch.where(
    padding_mask_bool,    # Condition (boolean mask where True means fill)
    min_dtype_val,        # Value if condition is True
    target_for_op         # Value if condition is False (original tensor slice)
)
causal_mask_clone[:, :, :, :mask_length] = filled_slice

With this change, the model generates text successfully on DirectML.

Entorno:

• Sistema Operativo: Windows 11 Pro Versión 24H2 (compilación de SO 26100.3915)
• Python: 3.12.3
• PyTorch: 2.2.1+cu121
• torch-directml: 0.2.1.dev240521
• transformers: 4.51.3
• NumPy: 1.26.4
• GPU: AMD Radeon RX 6700 XT
• Drivers de GPU: AMD Adrenalin 25.5.1 (Driver version from Task Manager: 32.0.21001.9024)

Script de Inferencia (inferencia.py):

import torch
import torch_directml as dml
from transformers import AutoTokenizer, AutoModelForCausalLM
import traceback

# Configuración del dispositivo DirectML
device = dml.device(0) 
print(f"Dispositivo configurado: {device}")

# Carga del modelo y el tokenizer
model_name = "EleutherAI/gpt-neo-1.3B"
print(f"Cargando el modelo {model_name} en {device}...")
try:
    model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
    print("Modelo cargado correctamente.")
except Exception as e:
    print(f"Error al cargar el modelo: {e}")
    traceback.print_exc()
    exit()

try:
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    print("Tokenizer configurado correctamente.")
except Exception as e:
    print(f"Error al cargar el tokenizer: {e}")
    traceback.print_exc()
    exit()

# Definir prompt
prompt = "La inteligencia artificial es"
print(f"Prompt de entrada: {prompt}")

inputs_on_device = {} 

try:
    print("Tokenizando el input...")
    inputs_cpu = tokenizer(prompt, return_tensors="pt", padding=True)
    
    inputs_on_device['input_ids'] = inputs_cpu['input_ids'].to(device=device, dtype=torch.int32)
    if 'attention_mask' in inputs_cpu:
        inputs_on_device['attention_mask'] = inputs_cpu['attention_mask'].to(device=device, dtype=torch.float32)
    else:
        inputs_on_device['attention_mask'] = torch.ones_like(inputs_on_device['input_ids'], dtype=torch.float32, device=device)
    
    print(f"Tipo de input_ids ANTES de la inferencia (en {device}): {inputs_on_device['input_ids'].dtype}")
    print(f"Tipo de attention_mask ANTES de la inferencia (en {device}): {inputs_on_device['attention_mask'].dtype}")

except Exception as e:
    print(f"Error durante la tokenización o movimiento a dispositivo: {e}")
    traceback.print_exc()
    exit()

try:
    print("Realizando inferencia...")
    with torch.no_grad(): 
        output = model.generate(
            inputs_on_device['input_ids'],
            attention_mask=inputs_on_device['attention_mask'],
            max_length=50,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    print(f"Output tensor (on device {output.device}): {output}")
    decoded_output = tokenizer.decode(output[0].to('cpu'), skip_special_tokens=True)
    print(f"Texto generado: {decoded_output}")
except Exception as e:
    print(f"Error durante la inferencia: {e}")
    traceback.print_exc()

Información Adicional:

• The aten::isin.Tensor_Tensor_out operator also shows a UserWarning for falling back to CPU, but this does not seem to be the cause of the critical uint8_t overflow error.
• Extensive environment debugging was performed to ensure correct PyTorch and torch-directml versions were loaded. The issue appears specific to masked_fill on the DML backend under these conditions with GPT-Neo.

---

This issue was debugged with the assistance of an AI. More context and discussions can be found
https://chatgpt.com/share/682100bf-4140-800e-837e-29b063d672c6
https://aistudio.google.com/app/prompts?state=%7B%22ids%22:%5B%22113S9a6wz6vZPeD5ZsK3gnvIJCoKbkmgb%22%5D,%22action%22:%22open%22,%22userId%22:%22108410992550650304901%22,%22resourceKeys%22:%7B%7D%7D&usp=sharing, https://drive.google.com/file/d/17iOuGsYs45WAMNC_GW2OxF_DdMUSosRZ/view?usp=sharing, https://drive.google.com/file/d/1MYQ6Sk1bU7H0ncMKgIuNe40FGHYom0Zo/view?usp=sharing, https://drive.google.com/file/d/1ZB2K--aFzLvDzydcnFeUGS55FG4wQgp8/view?usp=sharing, https://drive.google.com/file/d/1q7lb0IFjKrG2i-1lXnYYy426RZ3RzFXC/view?usp=sharing, https://drive.google.com/file/d/1sGHn5kqyWd7VuahTaVmnUbUit-M-ZEnI/view?usp=sharing, https://drive.google.com/file/d/1yGam3q17QlRkzxg5k5U2KIOPyTNvs4sd/view?usp=sharing

[jstoecker]

Thanks for reporting this, but DirectML is in maintenance mode. We will not be able to look into this issue, unfortunately.

[gucciwong]

Following up with the root cause, in case it's useful even with this closed as not_planned.

Traced the masked_fill/uint8_t overflow crash (and the torch.where workaround this issue suggests) further back. The workaround does eliminate the crash, but training then produces loss=nan from step 0 onward -- it only removes the symptom, not the underlying problem.

Per-module hooks on a Llama model traced the first NaN to layer 0's attention output, then to transformers' _unmask_unattended safety call being gated by attention_mask.device.type in ["cuda", "xpu"] (DirectML's device type is "privateuseone", so that call never runs). Adding it unconditionally surfaced that causal_mask was already a BoolTensor by that point -- not float. That traces to the original (unpatched) transformers code doing an in-place causal_mask *= bool_tensor a few lines earlier.

Minimal isolated repro confirms float_tensor *= bool_tensor (and the out-of-place form) silently degrades the entire result to bool dtype on DirectML, instead of following standard type promotion -- filed as a separate, more specific issue: #737.

The actual fix (verified with a full training run, 15 clean steps, loss decreasing, no NaN) doesn't even need the masked_fill->torch.where swap this issue suggests -- casting the bool comparison to the mask's dtype before the multiply resolves it at the root.