[BUG] `_ScaledDotProductAttention.forward` crashes under mixed precision (`16-mixed`) due to float16 overflow in `masked_fill`

## Description

`_ScaledDotProductAttention.forward` crashes with `RuntimeError: value cannot be converted to type at::Half without overflow` when using `TFTModel` with PyTorch Lightning's `precision="16-mixed"`.

The root cause is the hardcoded mask fill value `-1e9` in `tft_submodels.py`:

```python
attn = attn.masked_fill(mask, -1e9)
```

Under mixed precision, `torch.cuda.amp.autocast` runs `torch.bmm` in float16. The subsequent division by a CPU float32 scalar does **not** promote the tensor back to float32 under autocast. So `attn` remains float16 (`at::Half`) at the `masked_fill` call, and `-1e9` exceeds the float16 representable range (max ~65504), causing the overflow error.

A secondary issue: the scaling tensor is created on CPU with hardcoded float32 dtype, which is unnecessary cross-device/dtype overhead under autocast:

```python
dimension = torch.sqrt(torch.tensor(k.shape[-1]).to(torch.float32))
```

## Steps to Reproduce

```python
model = TFTModel(
    ...,
    pl_trainer_kwargs={
        "precision": "16-mixed",
        "accelerator": "gpu",
    },
)
model.fit(series=..., past_covariates=..., future_covariates=...)
```

## Proposed Fix

Two changes to `_ScaledDotProductAttention.forward`:

**1. Create the scaling tensor with matching dtype and device:**

```python
# Before:
dimension = torch.sqrt(torch.tensor(k.shape[-1]).to(torch.float32))

# After:
dimension = torch.sqrt(torch.tensor(k.shape[-1], dtype=attn.dtype, device=attn.device))
```

**2. Use dtype-safe fill value for masking:**

```python
# Before:
attn = attn.masked_fill(mask, -1e9)

# After:
attn = attn.masked_fill(mask, torch.finfo(attn.dtype).min)
```

Both changes are backward-compatible under float32. `torch.finfo(torch.float32).min ≈ -3.4e38` is functionally equivalent to `-1e9` for softmax masking.

## Environment

| | |
|---|---|
| Darts | 0.43.0 |
| PyTorch | 2.5.0 (CUDA 12.x) |
| PyTorch Lightning | 2.5.2 |
| Python | 3.10 |
| Precision | `"16-mixed"` |

## Full Stacktrace

<details>
<summary>Click to expand</summary>

```
2026-04-17 08:59:46,042 - GGE_TFT_Forecast - ERROR - An error occurred: value cannot be converted to type at::Half without overflow
Traceback (most recent call last):
  File "/.../gge_fc_tft.py", line 334, in <module>
    main()
  File "/.../gge_fc_tft.py", line 240, in main
    model.fit(
  File "/.../darts/utils/torch.py", line 95, in decorator
    return decorated(self, *args, **kwargs)
  File "/.../darts/models/forecasting/torch_forecasting_model.py", line 1031, in fit
    return self.fit_from_dataset(*params)
  File "/.../darts/utils/torch.py", line 95, in decorator
    return decorated(self, *args, **kwargs)
  File "/.../darts/models/forecasting/torch_forecasting_model.py", line 1220, in fit_from_dataset
    self._train(
  File "/.../darts/models/forecasting/torch_forecasting_model.py", line 1436, in _train
    trainer.fit(
  File "/.../pytorch_lightning/trainer/trainer.py", line 561, in fit
    call._call_and_handle_interrupt(
  File "/.../pytorch_lightning/trainer/call.py", line 48, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/.../pytorch_lightning/trainer/trainer.py", line 599, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/.../pytorch_lightning/trainer/trainer.py", line 1012, in _run
    results = self._run_stage()
  File "/.../pytorch_lightning/trainer/trainer.py", line 1056, in _run_stage
    self.fit_loop.run()
  File "/.../pytorch_lightning/loops/fit_loop.py", line 216, in run
    self.advance()
  File "/.../pytorch_lightning/loops/fit_loop.py", line 455, in advance
    self.epoch_loop.run(self._data_fetcher)
  File "/.../pytorch_lightning/loops/training_epoch_loop.py", line 152, in run
    self.advance(data_fetcher)
  File "/.../pytorch_lightning/loops/training_epoch_loop.py", line 344, in advance
    batch_output = self.automatic_optimization.run(trainer.optimizers[0], batch_idx, kwargs)
  File "/.../pytorch_lightning/loops/optimization/automatic.py", line 192, in run
    self._optimizer_step(batch_idx, closure)
  File "/.../pytorch_lightning/loops/optimization/automatic.py", line 270, in _optimizer_step
    call._call_lightning_module_hook(
  File "/.../pytorch_lightning/trainer/call.py", line 176, in _call_lightning_module_hook
    output = fn(*args, **kwargs)
  File "/.../pytorch_lightning/core/module.py", line 1328, in optimizer_step
    optimizer.step(closure=optimizer_closure)
  File "/.../pytorch_lightning/core/optimizer.py", line 154, in step
    step_output = self._strategy.optimizer_step(self._optimizer, closure, **kwargs)
  File "/.../pytorch_lightning/strategies/strategy.py", line 239, in optimizer_step
    return self.precision_plugin.optimizer_step(optimizer, model=model, closure=closure, **kwargs)
  File "/.../pytorch_lightning/plugins/precision/amp.py", line 79, in optimizer_step
    closure_result = closure()
  File "/.../pytorch_lightning/loops/optimization/automatic.py", line 146, in __call__
    self._result = self.closure(*args, **kwargs)
  File "/.../torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/.../pytorch_lightning/loops/optimization/automatic.py", line 131, in closure
    step_output = self._step_fn()
  File "/.../pytorch_lightning/loops/optimization/automatic.py", line 319, in _training_step
    training_step_output = call._call_strategy_hook(trainer, "training_step", *kwargs.values())
  File "/.../pytorch_lightning/trainer/call.py", line 328, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/.../pytorch_lightning/strategies/strategy.py", line 391, in training_step
    return self.lightning_module.training_step(*args, **kwargs)
  File "/.../darts/models/forecasting/pl_forecasting_module.py", line 256, in training_step
    return self._train_val_step(
  File "/.../darts/models/forecasting/pl_forecasting_module.py", line 292, in _train_val_step
    output = self._produce_train_output(
  File "/.../darts/models/forecasting/pl_forecasting_module.py", line 542, in _produce_train_output
    return self(self._process_input_batch(input_batch))
  File "/.../torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/.../torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/.../darts/models/forecasting/pl_forecasting_module.py", line 56, in forward_wrapper
    return forward(self, x_in, *args, **kwargs)
  File "/.../darts/models/forecasting/tft_model.py", line 617, in forward
    attn_out, attn_out_weights = self.multihead_attn(
  File "/.../torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/.../torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/.../darts/models/forecasting/tft_submodels.py", line 593, in forward
    head, attn = self.attention(qs, ks, vs, mask)
  File "/.../torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/.../torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/.../darts/models/forecasting/tft_submodels.py", line 549, in forward
    attn = attn.masked_fill(mask, -1e9)
RuntimeError: value cannot be converted to type at::Half without overflow
```

</details>

Navigation Menu

Uh oh!

[BUG] `_ScaledDotProductAttention.forward` crashes under mixed precision (`16-mixed`) due to float16 overflow in `masked_fill` #3081

Description

Steps to Reproduce

Proposed Fix

Environment

Full Stacktrace

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development


Darts	0.43.0
PyTorch	2.5.0 (CUDA 12.x)
PyTorch Lightning	2.5.2
Python	3.10
Precision	`"16-mixed"`

Uh oh!

[BUG] _ScaledDotProductAttention.forward crashes under mixed precision (16-mixed) due to float16 overflow in masked_fill #3081

Description

Description

Steps to Reproduce

Proposed Fix

Environment

Full Stacktrace

Activity

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

[BUG] `_ScaledDotProductAttention.forward` crashes under mixed precision (`16-mixed`) due to float16 overflow in `masked_fill` #3081