Skip to content

LagLlamaForecaster

LagLlamaForecaster

class LagLlamaForecaster(ckpt_path=None, device=None, context_length=32, num_samples=100, batch_size=1, use_rope_scaling=False, nonnegative_pred_samples=False, use_source_package=False, validation_split=0.2, trainer_kwargs=None, lr=0.0005, aug_prob=0.0)[source]

LagLlama Foundation Model for Time Series Forecasting.

LagLlama is a foundation model for univariate probabilistic time series forecasting based on a decoder-only transformer architecture. This implementation supports both zero-shot prediction using pretrained weights and fine-tuning on custom data.

The model checkpoint is automatically downloaded on first use if not provided.

Usage Workflows

This model supports two main workflows:

  • Zero-shot (default): Uses pretrained model as-is without training. Fast inference with no training overhead. Suitable for quick predictions. Simply call fit(y) then predict().

  • Fine-tuning via pretrain(): Fine-tune the model on panel/hierarchical data using the pretrain() method, then call fit(y) on a specific series to set context for prediction. This is the recommended workflow for domain-specific fine-tuning. Controlled by trainer_kwargs, lr, and aug_prob parameters.

Parameters:
ckpt_pathstr, optional (default=None)

Path to LagLlama checkpoint file. If None, automatically downloads from HuggingFace: “time-series-foundation-models/Lag-Llama”.

devicestr, optional (default=None)

Device for inference (“cpu”, “cuda”, “cuda:0”, etc.). If None, uses CUDA if available, otherwise CPU.

context_lengthint, optional (default=32)

Number of past time steps used as context for prediction. LagLlama was trained with context_length=32.

num_samplesint, optional (default=100)

Number of sample paths for probabilistic forecasting.

batch_sizeint, optional (default=1)

Batch size for prediction.

use_rope_scalingbool, optional (default=False)

Whether to use RoPE scaling for handling longer context lengths.

nonnegative_pred_samplesbool, optional (default=False)

If True, ensures all predicted samples are passed through ReLU.

use_source_packagebool, optional (default=False)

If True, uses the external lag-llama package instead of vendored version.

validation_splitfloat, optional (default=0.2)

Fraction of data for validation during pretrain(). Set to None to skip validation.

trainer_kwargsdict, optional (default=None)

Arguments passed to PyTorch Lightning Trainer during pretrain() (e.g., {"max_epochs": 10}). If None, defaults to {"max_epochs": 50}. Valid keys include:

acceleratorstr or Accelerator

Supports passing different accelerator types ("cpu", "gpu", "tpu", "hpu", "mps", "auto") as well as custom accelerator instances.

strategystr or Strategy, default=”auto”

Supports different training strategies with aliases as well as custom strategies.

deviceslist of int, str, or int, default=”auto”

The devices to use. Can be set to a positive number (int or str), a sequence of device indices (list or str), the value -1 to indicate all available devices should be used, or "auto" for automatic selection based on the chosen accelerator.

num_nodesint, default=1

Number of GPU nodes for distributed training.

precisionint, str, or None, default=”32-true”

Double precision (64, "64" or "64-true"), full precision (32, "32" or "32-true"), 16bit mixed precision (16, "16", "16-mixed") or bfloat16 mixed precision ("bf16", "bf16-mixed"). Can be used on CPU, GPU, TPUs, or HPUs.

loggerLogger, iterable of Logger, bool, or None, default=True

Logger (or iterable collection of loggers) for experiment tracking. A True value uses the default TensorBoardLogger if it is installed, otherwise CSVLogger. False will disable logging. If multiple loggers are provided, local files (checkpoints, profiler traces, etc.) are saved in the log_dir of the first logger.

callbackslist of Callback, Callback, or None, default=None

Add a callback or list of callbacks.

fast_dev_runint or bool, default=False

Runs n if set to n (int) else 1 if set to True batch(es) of train, val and test to find any bugs (i.e. a sort of unit test).

max_epochsint or None, default=None

Stop training once this number of epochs is reached. Disabled by default (None). If both max_epochs and max_steps are not specified, defaults to max_epochs = 1000. To enable infinite training, set max_epochs = -1.

min_epochsint or None, default=None

Force training for at least these many epochs. Disabled by default (None).

max_stepsint, default=-1

Stop training after this number of steps. Disabled by default (-1). If max_steps = -1 and max_epochs = None, will default to max_epochs = 1000. To enable infinite training, set max_epochs to -1.

min_stepsint or None, default=None

Force training for at least these number of steps. Disabled by default (None).

max_timestr, timedelta, dict of str to int, or None, default=None

Stop training after this amount of time has passed. Disabled by default (None). The time duration can be specified in the format DD:HH:MM:SS (days, hours, minutes, seconds), as a datetime.timedelta, or a dictionary with keys that will be passed to datetime.timedelta.

limit_train_batchesint, float, or None, default=1.0

How much of training dataset to check (float = fraction, int = num_batches). Value is per device.

limit_val_batchesint, float, or None, default=1.0

How much of validation dataset to check (float = fraction, int = num_batches). Value is per device.

limit_test_batchesint, float, or None, default=1.0

How much of test dataset to check (float = fraction, int = num_batches). Value is per device.

limit_predict_batchesint, float, or None, default=1.0

How much of prediction dataset to check (float = fraction, int = num_batches). Value is per device.

overfit_batchesint or float, default=0.0

Overfit a fraction of training/validation data (float) or a set number of batches (int).

val_check_intervalint, float, str, timedelta, dict, or None, default=1.0

How often to check the validation set. Pass a float in the range [0.0, 1.0] to check after a fraction of the training epoch. Pass an int to check after a fixed number of training batches. An int value can only be higher than the number of training batches when check_val_every_n_epoch=None, which validates after every N training batches across epochs or during iteration-based training. Additionally, accepts a time-based duration as a string "DD:HH:MM:SS", a datetime.timedelta, or a dict of kwargs to datetime.timedelta. When time-based, validation triggers once the elapsed wall-clock time since the last validation exceeds the interval; the check occurs after the current batch completes, the validation loop runs, and the timer is reset.

check_val_every_n_epochint or None, default=1

Perform a validation loop after every N training epochs. If None, validation will be done solely based on the number of training batches, requiring val_check_interval to be an integer value. When used together with a time-based val_check_interval and check_val_every_n_epoch > 1, validation is aligned to epoch multiples: if the interval elapses before the next multiple-N epoch, validation runs at the start of that epoch (after the first batch) and the timer resets; if it elapses during a multiple-N epoch, validation runs after the current batch. For None or 1 cases, the time-based behavior of val_check_interval applies without additional alignment.

num_sanity_val_stepsint or None, default=2

Sanity check runs n validation batches before starting the training routine. Set it to -1 to run all batches in all validation dataloaders.

log_every_n_stepsint or None, default=50

How often to log within steps.

enable_checkpointingbool or None, default=True

If True, enable checkpointing. It will configure a default ModelCheckpoint callback if there is no user-defined ModelCheckpoint in callbacks.

enable_progress_barbool or None, default=True

Whether to enable the progress bar by default.

enable_model_summarybool or None, default=True

Whether to enable model summarization by default.

accumulate_grad_batchesint, default=1

Accumulates gradients over k batches before stepping the optimizer.

gradient_clip_valint, float, or None, default=None

The value at which to clip gradients. Passing gradient_clip_val=None disables gradient clipping. If using Automatic Mixed Precision (AMP), the gradients will be unscaled before.

gradient_clip_algorithmstr or None, default=None

The gradient clipping algorithm to use. Pass gradient_clip_algorithm="value" to clip by value, and gradient_clip_algorithm="norm" to clip by norm. By default it will be set to "norm".

deterministicbool, {“warn”}, or None, default=None

If True, sets whether PyTorch operations must use deterministic algorithms. Set to "warn" to use deterministic algorithms whenever possible, throwing warnings on operations that don’t support deterministic mode. If not set, defaults to False.

benchmarkbool or None, default=None

The value (True or False) to set torch.backends.cudnn.benchmark to. The value for torch.backends.cudnn.benchmark set in the current session will be used (False if not manually set). If deterministic is set to True, this will default to False. Override to manually set a different value.

inference_modebool

Whether to use torch.inference_mode() or torch.no_grad() during evaluation (validate/test/predict).

use_distributed_samplerbool

Whether to wrap the DataLoader’s sampler with torch.utils.data.DistributedSampler. If not specified this is toggled automatically for strategies that require it. By default, it will add shuffle=True for the train sampler and shuffle=False for validation/test/predict samplers. If you want to disable this logic, you can pass False and add your own distributed sampler in the dataloader hooks. If True and a distributed sampler was already added, Lightning will not replace the existing one. For iterable-style datasets, this is not done automatically.

profilerProfiler, str, or None, default=None

To profile individual steps during training and assist in identifying bottlenecks.

detect_anomalybool, default=False

Enable anomaly detection for the autograd engine.

barebonesbool

Whether to run in “barebones mode”, where all features that may impact raw speed are disabled. This is meant for analyzing the Trainer overhead and is discouraged during regular training runs. The following features are deactivated: enable_checkpointing, logger, enable_progress_bar, log_every_n_steps, enable_model_summary, num_sanity_val_steps, fast_dev_run, detect_anomaly, profiler, log(), log_dict().

pluginslist, object, or None, default=None

Precision, ClusterEnvironment, CheckpointIO, LayerSync, a list of those, or None. Plugins allow modification of core behavior like ddp and amp, and enable custom lightning plugins.

sync_batchnormbool, default=False

Synchronize batch norm layers between process groups/whole world.

reload_dataloaders_every_n_epochsint, default=0

Set to a positive integer to reload dataloaders every n epochs.

default_root_dirstr, Path, or None, default=os.getcwd()

Default path for logs and weights when no logger/ckpt_callback passed. Can be remote file paths such as s3://mybucket/path or hdfs://path/.

enable_autolog_hparamsbool, default=True

Whether to log hyperparameters at the start of a run.

model_registrystr or None

The name of the model being uploaded to Model hub.

lrfloat, optional (default=5e-4)

Learning rate for fine-tuning during pretrain().

aug_probfloat, optional (default=0.0)

Data augmentation probability during pretrain().

Attributes:
cutoff

Cut-off = “present time” state of forecaster.

fh

Forecasting horizon that was passed.

is_fitted

Whether fit has been called.

state

State of the estimator.

References

[1]

Rasul, Kashif, et al. “Lag-Llama: Towards Foundation Models for Probabilistic Time Series Forecasting.” arXiv preprint arXiv:2310.08278 (2023).

Examples

Zero-shot forecasting (default)

>>> from sktime.forecasting.lagllama import LagLlamaForecaster
>>> from sktime.forecasting.base import ForecastingHorizon
>>> from sktime.datasets import load_airline
>>>
>>> y = load_airline()
>>> forecaster = LagLlamaForecaster(
...     context_length=32,
...     num_samples=100
... )
>>> fh = ForecastingHorizon([1, 2, 3, 4, 5, 6])
>>> forecaster.fit(y, fh=fh)
LagLlamaForecaster(...)
>>> y_pred = forecaster.predict()  # Point predictions
>>> # 90% prediction intervals
>>> y_interval = forecaster.predict_interval(coverage=0.9)

Fine-tuning with pretrain() on panel data

>>> from sktime.forecasting.lagllama import LagLlamaForecaster
>>> from sktime.datasets import load_airline
>>> from sktime.utils._testing.hierarchical import (
...     _make_hierarchical,
... )
>>>
>>> # Create panel data for pretraining
>>> y_panel = _make_hierarchical(
...     hierarchy_levels=(3,), min_timepoints=50, max_timepoints=50
... )
>>> # Fine-tune on panel data
>>> forecaster = LagLlamaForecaster(
...     context_length=32,
...     num_samples=100,
...     trainer_kwargs={"max_epochs": 5},
...     lr=5e-4,
...     validation_split=0.2
... )
>>> forecaster.pretrain(y_panel)  # Fine-tune on panel
LagLlamaForecaster(...)
>>> # Now fit to specific series and predict
>>> y = load_airline()
>>> forecaster.fit(y, fh=[1, 2, 3, 4, 5, 6])
LagLlamaForecaster(...)
>>> y_pred = forecaster.predict()

Methods

check_is_fitted([method_name])

Check if the estimator has been fitted.

check_range_index(df)

Check if the index is a range index.

clone()

Obtain a clone of the object with same hyper-parameters and config.

clone_tags(estimator[, tag_names])

Clone tags from another object as dynamic override.

create_test_instance([parameter_set])

Construct an instance of the class, using first test parameter set.

create_test_instances_and_names([parameter_set])

Create list of all test instances and a list of names for them.

fit(y[, X, fh])

Fit forecaster to training data.

fit_predict(y[, X, fh, X_pred])

Fit and forecast time series at future horizon.

get_class_tag(tag_name[, tag_value_default])

Get class tag value from class, with tag level inheritance from parents.

get_class_tags()

Get class tags from class, with tag level inheritance from parent classes.

get_config()

Get config flags for self.

get_fitted_params([deep])

Get fitted parameters.

get_param_defaults()

Get object's parameter defaults.

get_param_names([sort])

Get object's parameter names.

get_params([deep])

Get a dict of parameters values for this object.

get_pretrained_params([deep])

Get pretrained parameters of this estimator.

get_tag(tag_name[, tag_value_default, ...])

Get tag value from instance, with tag level inheritance and overrides.

get_tags()

Get tags from instance, with tag level inheritance and overrides.

get_test_params([parameter_set])

Return testing parameter settings for the estimator.

handle_range_index(index)

Convert RangeIndex to Dummy DatetimeIndex.

infer_freq(index)

Infer frequency of the index.

is_composite()

Check if the object is composed of other BaseObjects.

load_from_path(serial)

Load object from file location.

load_from_serial(serial)

Load object from serialized memory container.

predict([fh, X])

Forecast time series at future horizon.

predict_interval([fh, X, coverage])

Compute/return prediction interval forecasts.

predict_proba([fh, X, marginal])

Compute/return fully probabilistic forecasts.

predict_quantiles([fh, X, alpha])

Compute/return quantile forecasts.

predict_residuals([y, X])

Return residuals of time series forecasts.

predict_var([fh, X, cov])

Compute/return variance forecasts.

pretrain(y[, X, fh])

Pre-train forecaster on panel (global) data.

reset()

Reset the object to a clean post-init state.

return_time_index(df)

Return the time index, given any type of index.

save([path, serialization_format])

Save serialized self to bytes-like object or to (.zip) file.

score(y[, X, fh])

Scores forecast against ground truth, using MAPE (non-symmetric).

set_config(**config_dict)

Set config flags to given values.

set_params(**params)

Set the parameters of this object.

set_random_state([random_state, deep, ...])

Set random_state pseudo-random seed parameters for self.

set_tags(**tag_dict)

Set instance level tag overrides to given values.

update(y[, X, update_params])

Update cutoff value and, optionally, fitted parameters.

update_predict(y[, cv, X, update_params, ...])

Make predictions and update model iteratively over the test set.

update_predict_single([y, fh, X, update_params])

Update model with new data and make forecasts.