[BUG] Inference-mode attention mask in TFT causing RuntimeError

[daidahao]

Describe the bug
A RuntimeError could be thrown when training the TFT again after forecasting.

To Reproduce
Steps to reproduce the behavior, preferably code snippet.

import numpy as np

from darts import TimeSeries
from darts.models import TFTModel

# Create a series with 3 components
series = TimeSeries.from_values(np.random.rand(100, 3).astype(np.float32))

# Define a TFT model
model = TFTModel(10, 10, add_relative_index=True, n_epochs=1)

# Train the model
model.fit(series, verbose=True)

# Make a prediction with `num_samples` equal to the batch size
model.predict(10, num_samples=model.batch_size)

# Train the model again, causing a `RuntimeError` due to uncleaned `attention_mask`
model.fit(series, verbose=True)

A RuntimeError is thrown because attention_mask from inference is being reused for retraining.

RuntimeError: Inference tensors cannot be saved for backward. To work around you can make a clone to get a normal tensor and use it in autograd.

Expected behavior
A clear and concise description of what you expected to happen.

System (please complete the following information):

• Python version: 3.13
• darts version: latest from source

Additional context
The bug should be easily fixable by clearing the attention_mask in TFT before each training.

[dennisbader]

Hi @daidahao and thanks for raising this issue 🚀

Yes, we should fix this as it's a bug. On a side note: we discourage calling model.fit() multiple times for torch models, and instead recommend the following section for fine-tuning.

If you want to train your model from scratch, you should use something like:

model = model.untrained_model()
model.fit()

Would you like to contribute with the fix? :)

[daidahao]

@dennisbader

Thank you for your great advice on fine-tuning!

However, please correct me if I am wrong: every time fit() is called on a torch model, an optimizer would be configured. Even though optimizer_cls and optimizer_kwargs would be the same from model declaration, the optimizer would be new because of Lightning trainer fit() behaviors. In that case, there should not be any issues with using the same model instance for fine-tuning.

Also, from Darts v0.38.0, there is a load_best option for fit(), thanks to #2903. Using the new option, creating an identical model and calling load_weights_from_checkpoint() can be replaced. With that, the following script would still cause a RuntimeError:

import numpy as np

from darts import TimeSeries
from darts.models import TFTModel

# Create a series with 3 components
series = TimeSeries.from_values(np.random.rand(100, 3).astype(np.float32))
train_series, val_series = series.split_before(0.8)

# Define a TFT model
model = TFTModel(
    input_chunk_length=10,
    output_chunk_length=10,
    add_relative_index=True,
    n_epochs=3,
    save_checkpoints=True,
    pl_trainer_kwargs={
        "accelerator": "cpu",
    }
)

# Train the model
model.fit(
    series=train_series,
    val_series=val_series,
    verbose=True,
    load_best=True,
)

# Make a prediction with `num_samples` equal to the batch size
model.predict(10, num_samples=model.batch_size)

# Train the model again, causing a `RuntimeError` due to uncleaned `attention_mask`
model.fit(series, verbose=True)

Would you like to contribute with the fix? :)

Yes, I will contribute a fix to this bug.

[dennisbader]

Hi @daidahao.

However, please correct me if I am wrong: every time fit() is called on a torch model, an optimizer would be configured. Even though optimizer_cls and optimizer_kwargs would be the same from model declaration, the optimizer would be new because of Lightning trainer fit() behaviors. In that case, there should not be any issues with using the same model instance for fine-tuning.

Yes, it would use a new optimizer, but the model weights / params are from the latest state after training. Usually for fine-tuning would see different optimizer kwargs (e.g. lower learning rates, ...) and learning rate schedulers to "fine" tune on new data.

Also, from Darts v0.38.0, there is a load_best option for fit(), thanks to #2903. Using the new option, creating an identical model and calling load_weights_from_checkpoint() can be replaced.

Yes, darts 0.38.0 added load_best, which will load the best model (but not only the weights). It would be similar to load_from_checkpoint which will have the same optimizer kwargs and so on (that's different to load_weights_from_checkpoint where you attach the weights to a new model you created with different optimizer kwargs, ...).

With that, the following script would still cause a RuntimeError:

Yes, it's because you call predict() after fit(), and then try to fit the TFTModel again which is running into the bug :)

If you load the checkpoint after prediction it would work.

# Train the model
model.fit(
    series=train_series,
    val_series=val_series,
    verbose=True,
)

# Make a prediction with `num_samples` equal to the batch size
model.predict(
    10,
    num_samples=model.batch_size
)

# load from checkpoint
model = TFTModel.load_from_checkpoint(
    model_name=model.model_name,
    work_dir=model.work_dir,
    best=True
)

# Train the model again, causing a `RuntimeError` due to uncleaned `attention_mask`
model.fit(series, verbose=True)

Btw, something that I noticed is that this bug only occurs when setting num_samples>1 in model.predict()

[daidahao]

Hi @dennisbader

Yes, it would use a new optimizer, but the model weights / params are from the latest state after training. Usually for fine-tuning would see different optimizer kwargs (e.g. lower learning rates, ...) and learning rate schedulers to "fine" tune on new data.

Thank you for the explanations. Yes, that was a special use case of mine where I want to fine-tune the model using the same learning rate. Like you said, it is best to fine-tune the model with a different learning rate, scheduler, etc.

Btw, something that I noticed is that this bug only occurs when setting num_samples>1 in model.predict()

Yes, that is because the bug is only triggered when num_samples is equal to the last batch size. Otherwise, the attention mask would be regenerated.

I wonder now if there is still need to fix the bug since the user is advised to load_from_checkpoint(). The bug only occurs (as far as I know) when fit() is called after predict(). What do you think?