LoadTime

English | 日本語

LoadTime is a Python package designed specifically to tackle the challenge of long waiting times associated with loading large-scale pretrained language models, such as HuggingFace models, into GPU or CPU memory.

With LoadTime, instead of waiting in uncertainty, you can visualize the progress of your loading process.

Of course, it can also be used for other long-term operations.

Installation

You can install LoadTime via pip:

pip install loadtime

Key Features

• Real-time tracking: LoadTime provides real-time tracking of the loading process. No more staring at a static screen!

• Progress Bar: The package displays a progress bar, showing you how much of the process has been completed and how much is still remaining. It takes the guesswork out of waiting!

• Past Loading Time Cache: One unique feature of LoadTime is its ability to remember the time it took to load a model in the past. The package automatically caches the total loading time of your operations. The next time you load the same model, LoadTime uses this cached information to provide an even more accurate progress bar.

• Customizable Display: LoadTime allows you to customize the progress display with your own messages. You can tailor the tool to fit your personal needs.

• Optimized for HuggingFace Models: LoadTime has been optimized for loading HuggingFace models, with special handling of the download progress display when the model is not cached locally.

Basic Usage

Here is a simple example of how to use the LoadTime package:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from loadtime import LoadTime

model_path = "togethercomputer/RedPajama-INCITE-Chat-3B-v1"

model = LoadTime(name=model_path,
                 fn=lambda: AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16))()

tokenizer = AutoTokenizer.from_pretrained(model_path) # important: load tokenizer after loading model

Initial Parameters

Take control of your loading times with LoadTime!