Warm Start Training: Initialize from Existing Models

What is the warm start feature and how to use it?

Warm start initializes a new training job with weights from an already-trained model, so you can build on previous work instead of starting from scratch. Pass the ID of an existing training job to Trainer.fit():

1
from kumoai.trainer import Trainer
2
3
original_job_id = "original_training_job_id"
4
trainer = Trainer(model_plan=model_plan)
5
trainer.fit(
6
    graph=graph,
7
    train_table=train_table,
8
    warm_start_job_id=original_job_id,
9
)

The trainer loads the final model weights and optimizer state from original_job_id. It inherits the original job’s model architecture, column processing, and neighbor sampling configurations. Inheriting these settings prevents you from making architecture changes that would be incompatible with the loaded weights.

You can customize the training job and optimization plans when warm starting. Refer to the Trainer.fit() reference and Model Plan.

Benefits of warm start

• Stable embeddings: Update embeddings over time without causing a drastic shift in the embedding space. This helps maintain consistency for downstream tasks such as retrieval or ranking. Achieving stable embeddings may require tuning optimization parameters such as the learning rate (lr) and number of epochs (max_epochs).
• Faster retraining: Warm start is especially useful when data distribution (such as user behavior) changes dynamically over time and you need to retrain frequently to keep the model current.

Requirements and Limitations

• The warm start job must use the same task type as the original job.
• Forecasting: Not supported.
• Classification: Target classes must match the original job exactly.