Faster Cold Starts

tensorizer = ">=2.7.0"

from tensorizer import TensorSerializer
def serialize_model(model, save_path):
    """Serialize the model and save the weights to the save_path."""
    try:
        serializer = TensorSerializer(save_path)
        start = time.time()
        serializer.write_module(model)
        end = time.time()
        print(f"Serializing model took {end - start} seconds", file=sys.stderr)
        serializer.close()
        return True
    except Exception as e:
        print("Serialization failed with error:", e, file=sys.stderr)
        return False

from tensorizer import TensorDeserializer
from tensorizer.utils import no_init_or_tensor
def deserialize_saved_model(model_path, model_id, plaid=True):
    """Deserialize the model from the model_path and load into GPU memory."""

    # create a config object that we can use to init an empty model
    config = AutoConfig.from_pretrained(model_id)

    # Initialize empty model without loading weights into GPU
    print("Initializing empty model", file=sys.stderr)
    start = time.time()
    with no_init_or_tensor():
        # Load empty model from config
        model = AutoModelForCausalLM.from_config(config)
    end_init = time.time() - start

    # Create deserializer object
    # Note: plaid_mode enables faster deserialization but isn't safe for training
    deserializer = TensorDeserializer(model_path, plaid_mode=True)

    # Deserialize model directly into GPU (zero-copy)
    print("Loading model", file=sys.stderr)
    start = time.time()
    deserializer.load_into_module(model)
    end = time.time()
    deserializer.close()

    # Report timings
    print(f"Initializing empty model took {end_init} seconds", file=sys.stderr)
    print(f"\nDeserializing model took {end - start} seconds\n", file=sys.stderr)

    return model