Chronos Forecasting

You can train Chronos models from scratch or fine-tune existing checkpoints.

Configuration

Modify the training configs in training/configs. Key keys include:

• training_data_paths: A list of paths to your Arrow files.
• probability: A list of mixing probabilities for each dataset file.

To fine-tune a pretrained model, set model_id to the checkpoint (e.g., amazon/chronos-t5-small), set random_init: false, and adjust max_steps and learning_rate.

Execution

Run the training script using python training/train.py with a --config flag.

Tips:

• If training is slow, try setting torch_compile to false or adjusting shuffle_buffer_length.
• Important: When pretraining causal models (like GPT2), the script uses LastValueImputation for missing values. Ensure your input data is imputed similarly for consistent results in ChronosPipeline.predict().

Outputs and checkpoints are saved in output/run-{id}/.

# On single GPU
CUDA_VISIBLE_DEVICES=0 python training/train.py --config /path/to/modified/config.yaml

# On multiple GPUs (example with 8 GPUs)
torchrun --nproc-per-node=8 training/train.py --config /path/to/modified/config.yaml

# Fine-tune amazon/chronos-t5-small for 1000 steps
CUDA_VISIBLE_DEVICES=0 python training/train.py --config /path/to/modified/config.yaml \
    --model-id amazon/chronos-t5-small \
    --no-random-init \
    --max-steps 1000 \
    --learning-rate 0.001

Evaluate models using the evaluation/evaluate.py script to compute WQL and MASE values for in-domain or zero-shot benchmarks.

Running Evaluation

Use the --chronos-model-id to specify the model, --batch-size, --device, and --num-samples.

Aggregating Scores

You can use the following logic to compute aggregated relative WQL and MASE scores against a baseline using scipy.stats.gmean.

# In-domain evaluation
python evaluation/evaluate.py evaluation/configs/in-domain.yaml evaluation/results/chronos-t5-small-in-domain.csv \
    --chronos-model-id "amazon/chronos-t5-small" \
    --batch-size=32 \
    --device=cuda:0 \
    --num-samples 20

# Zero-shot evaluation
python evaluation/evaluate.py evaluation/configs/zero-shot.yaml evaluation/results/chronos-t5-small-zero-shot.csv \
    --chronos-model-id "amazon/chronos-t5-small" \
    --batch-size=32 \
    --device=cuda:0 \
    --num-samples 20

import pandas as pd
from scipy.stats import gmean


def agg_relative_score(model_df: pd.DataFrame, baseline_df: pd.DataFrame):
    relative_score = model_df.drop("model", axis="columns") / baseline_df.drop(
        "model", axis="columns"
    )
    return relative_score.agg(gmean)


result_df = pd.read_csv("evaluation/results/chronos-t5-small-in-domain.csv").set_index("dataset")
baseline_df = pd.read_csv("evaluation/results/seasonal-naive-in-domain.csv").set_index("dataset")

agg_score_df = agg_relative_score(result_df, baseline_df)

Chronos-2 supports cross-learning via the cross_learning=True parameter. This allows the model to share information across all time series in a batch, which is particularly beneficial for forecasting multiple related time series with short historical context.

Important Considerations:

• Task-dependent results: Cross-learning may not always improve forecasts; evaluate it for your specific use case.
• Batch size dependency: Results depend on batch_size. For optimal results, consider a batch size around 100.
• Input homogeneity: Works best with homogeneous inputs (e.g., multiple univariate time series of the same type).
• Short context benefit: Most helpful when individual series have limited historical context.

# Enable cross-learning for joint prediction
joint_pred_df = pipeline.predict_df(
    context_df,
    prediction_length=24,
    quantile_levels=[0.1, 0.5, 0.9],
    cross_learning=True,
    batch_size=100,
)

Batch Transform is used for large-scale offline processing. Data must be staged in S3 in JSONL format, where each line is a valid Chronos API payload.

Workflow:

1. Prepare input data: Convert DataFrames to JSON payloads and split into chunks (e.g., 100 series per line).
2. Upload JSONL files to S3.
3. Use sagemaker.transformer.Transformer to run the job.

Configuration:

• strategy="SingleRecord": Processes one JSON line at a time.
• assemble_with="Line": Combines results line by line.
• accept="application/json".

from sagemaker.transformer import Transformer

output_s3_uri = f"s3://{bucket}/{s3_prefix}/batch-output/"

transformer = Transformer(
    model_name=chronos_model.name,
    instance_count=1,
    instance_type="ml.c5.4xlarge",
    output_path=output_s3_uri,
    strategy="SingleRecord",
    assemble_with="Line",
    accept="application/json",
)

transformer.transform(
    data=input_s3_uri,
    content_type="application/json",
    split_type="Line",
    wait=True,
)