Tao Finetune Kumo Forecast

Kumo-Forecast

Transformer-based multivariate time series forecasting using self-supervised pretraining on diverse temporal data. Three inference modes: standard (direct forecast), DARR (context-enhanced kNN retrieval blending), and interpretability (latent trajectory extraction, semantic flow, lag×horizon attribution, trajectory stability, and diagnostic ratios — full explanation bundle with PDF report). Fine-tuning adapts the forecasting head — and optionally the cross-channel layer — to your domain.

Source code: https://github.com/NVIDIA/Kumo-TS Pretrained weights: https://huggingface.co/nvidia/Kumo-Forecast

External dependencies

Credentials

nvidia/Kumo-Forecast is a public repo — no token required for downloading weights. If you hit a 401/403 (gated access or license not accepted) or a 504 on first download, see the Known pitfalls section.

Quick start

bashCopy code
git clone --branch main --single-branch https://github.com/NVIDIA/Kumo-TS Kumo-TS
cd Kumo-TS/Kumo-Forecast
uv sync --group dev
uv pip install -e .          # editable install — required for clean sdk.* imports

# Standard inference (auto-downloads weights from HF on first run, no auth needed)
uv run python sdk/quick_example.py

Inference

Import and call perform_forecasting from sdk/forecasting.py. It auto-downloads weights, standardizes input, runs autoregressive rollout for long horizons, and returns a DataFrame with {target_column}_forecast rows for the requested horizon.

pythonCopy code
import sys, pandas as pd
sys.path.append("/path/to/Kumo-TS/Kumo-Forecast")  # clone Kumo-TS with --branch main
from sdk.forecasting import ForecastingConfig, perform_forecasting

df = pd.read_csv("your_data.csv")   # must have timestamp + numeric target column

config = ForecastingConfig(
    timestamp_column="timestamp",    # parseable datetime column
    target_column="target",          # primary target to forecast
    seq_len=512,                     # input context length (rows consumed)
    forecast_horizon=72,             # steps ahead to predict (max 512)
    model_horizon=72,                # native model horizon; change when using custom weights
    standardizer_pkl="standardizer.pkl",   # auto-downloaded from HF if missing
    ckpt="run8_best_model_cr.pt",          # auto-downloaded; see Checkpoints table
)
results = perform_forecasting(df=df, config=config)
# Returns DataFrame: timestamp | {target_column}_forecast  (forecast_horizon rows)
print(results.head())

For regular inference, put forecasting arguments in YAML and pass the file as config. Start from the fully commented template in sdk/forecasting_inference_config.yaml.

yamlCopy code
inference:
  # Number of historical rows consumed for the input context window.
  seq_len: 512
  # Number of future rows to predict.
  forecast_horizon: 72
  # Native horizon used when training the checkpoint.
  model_horizon: 72
  # Enable cross-channel attention for multivariate forecasting.
  use_cross_channel: true

pythonCopy code
results = perform_forecasting(df=df, config="sdk/forecasting_inference_config.yaml")

The YAML can be flat or use a top-level inference: section inside a larger TAO/AutoML spec. Unknown keys raise ValueError.

Checkpoints

Set use_cross_channel=False in ForecastingConfig or YAML to use the standard checkpoint:

pythonCopy code
results = perform_forecasting(df=df, config=ForecastingConfig(use_cross_channel=False))

DARR mode (context-enhanced forecasting)

Supply context_df to enable DARR: the SDK builds a kNN memory from historical windows and blends direct predictions with retrieved neighbors (alpha * direct + (1 - alpha) * kNN).

pythonCopy code
context_df = pd.read_csv("historical_data.csv")   # needs ≥ seq_len + model_horizon rows

config = ForecastingConfig(
    forecast_horizon=72,
    alpha=0.2,                  # 0.2 = 20% direct, 80% kNN (default: 0.01)
    k=64,                       # number of nearest neighbors
    temperature=0.05,           # kNN softmax temperature
)
results = perform_forecasting(df=df, config=config, context_df=context_df)

Context and input datasets do not need identical columns — the SDK aligns to common features and warns when columns differ. Both must share timestamp_column and target_column.

Interpretability

Set interpretability=True to activate the Model-Agnostic Interpretability Framework. It produces localized, horizon-specific, time-aware explanations — including lag×horizon attribution, semantic flow, trajectory stability, diagnostic ratios, and (for multivariate inputs) channel-axis attribution and coupling analysis.

For the full parameter reference, output bundle, and component descriptions, see Kumo-Forecast/README.md.

Fine-tuning

Fine-tune the forecasting head (encoder/embedder frozen by default) on your own time series. --ckpt-init auto warm-starts from the published Kumo-Forecast checkpoint; --ckpt-init none trains a fresh head from the base backbone.

bashCopy code
cd /path/to/Kumo-TS/Kumo-Forecast
# Without cross-channel (uses moment_head_512_6hr.pt)
uv run python examples/finetune_example.py \
  --csv /path/to/timeseries.csv \
  --timestamp-col timestamp \
  --target-cols target \
  --seq-len 512 --forecast-horizon 72 \
  --epochs 5 --batch-size 8 --lr 1e-4 \
  --output-dir artifacts/finetune_my_data

# With cross-channel layer (uses run8_best_model_cr.pt)
uv run python examples/finetune_example.py \
  --csv /path/to/timeseries.csv \
  --timestamp-col timestamp \
  --target-cols sensor_1,sensor_2,sensor_3 \
  --use-cross-channel --cross-channel-heads 8 \
  --epochs 5 \
  --output-dir artifacts/finetune_cross_channel

Fine-tuning arguments

*One of --csv or --train-csv is required.

Inference with fine-tuned checkpoint

pythonCopy code
config = ForecastingConfig(
    timestamp_column="timestamp",
    target_column="target",
    seq_len=512,
    forecast_horizon=72,
    model_horizon=72,
    standardizer_pkl="artifacts/finetune_my_data/standardizer.pkl",
    ckpt="artifacts/finetune_my_data/best_model.pt",
    use_cross_channel=False,   # set True if trained with --use-cross-channel
)
results = perform_forecasting(df=df, config=config)

Data requirements

Output structure

Inference (standard / DARR):

Copy code
DataFrame: timestamp | {target_column}_forecast   (forecast_horizon rows)

Fine-tuning (--output-dir artifacts/finetune_my_data):

Copy code
artifacts/finetune_my_data/
├── best_model.pt            # checkpoint with lowest validation MSE
├── standardizer.pkl         # normalization statistics for this dataset
├── finetune_metadata.json   # model config, channels, best epoch, all args
├── metrics.json             # scalar summary: {"val_mse": float, "val_mae": float} — consumed by AutoML runner
└── epoch_metrics.json       # per-epoch list: [{epoch, train_mse, val_mse, val_mae}, ...]

Hardware

AutoML (HPO: hyperparameter optimization)

This skill is AutoML-enabled for both fine-tuning and DARR inference. When an HPO request arrives, route it through tao-skill-bank:tao-run-automl with this model's skill_dir.

Read references/automl.md when the user asks for AutoML/HPO setup, tunable parameters, runner examples, inference trial scripts, DARR HPO, or AutoML result handoff details.

Known pitfalls