torch-directml: Tensor *= bool_tensor silently produces bool dtype instead of following standard type promotion

[gucciwong]

Summary

torch.Tensor.mul_/*= (and the out-of-place *) with a bool right-hand operand does not follow PyTorch's standard type promotion rules on the DirectML (privateuseone) backend. On CPU/CUDA, float_tensor *= bool_tensor promotes the bool to the float dtype and keeps a float result. On DirectML, the entire result is silently degraded to bool dtype, discarding the float values.

Minimal repro

import torch
import torch_directml

dev = torch_directml.device()
x = torch.full((4,), -3.4e38, dtype=torch.float32, device=dev)
print("dtype before:", x.dtype)          # torch.float32

bool_mask = torch.tensor([True, False, True, False], device=dev)
x *= bool_mask
print("dtype after *= bool:", x.dtype)   # torch.bool  <-- expected torch.float32
print(x)                                  # tensor([True, False, True, False]) -- original float values lost

Same result for the out-of-place form and .mul_() explicitly:

x1 = torch.full((4,), -3.4e38, dtype=torch.float32, device=dev)
y1 = x1 * bool_mask
print(y1.dtype)  # torch.bool

x2 = torch.full((4,), -3.4e38, dtype=torch.float32, device=dev)
x2.mul_(bool_mask)
print(x2.dtype)  # torch.bool

Casting the bool tensor to the target dtype before the multiply avoids the issue entirely and produces the correct result:

x3 = torch.full((4,), -3.4e38, dtype=torch.float32, device=dev)
x3 *= bool_mask.float()
print(x3.dtype, x3)  # torch.float32, values preserved correctly

Environment

• torch: 2.4.1+cpu
• torch-directml: 0.2.5.dev240914
• GPU: AMD Radeon RX 7600M XT (gfx1102)
• OS: Windows-11-10.0.26200-SP0

Why this matters / how it was found

Found while investigating why fine-tuning a Llama model with transformers+peft on torch-directml crashes with:

RuntimeError: value cannot be converted to type uint8_t without overflow

transformers' causal-mask construction (_prepare_4d_causal_attention_mask_with_cache_position in modeling_llama.py, and the same pattern is duplicated across most model files in the library) does:

causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)

This is exactly float_tensor *= bool_tensor. On DirectML, causal_mask silently becomes a BoolTensor at this line, before it ever reaches masked_fill/torch.where further down — which is what actually causes the crash reported in #702, and (after patching around that crash by swapping masked_fill for torch.where, per #702's own suggested workaround) also explains why training then produces loss=nan from step 0 onward instead of actually working: causal_mask was never a valid float mask to begin with, torch.where doesn't fix that upstream corruption.

Casting the bool comparison to the target dtype before the multiply ((torch.arange(...) > cache_position...).to(dtype)) fixes this at the root — verified with a full LoRA fine-tuning run: 15 clean training steps, loss decreasing smoothly, no crash, no NaN, with no change to masked_fill at all.

This is likely a more fundamental root cause than #702 describes, since it's a general type-promotion bug in the multiply kernel rather than something specific to masked_fill's handling of extreme fill values. Any DirectML code path doing float_tensor *= bool_tensor (in-place or not) would hit this, not just this one masking function.

Full writeup with more findings on this hardware/backend combination: https://github.com/gucciwong/amd-local-ai-bench/blob/main/docs/training-methodologies.md