[ChatStream] Adding a Progress Bar to Slow Model Loading
Hello, this is the Product Development Division at Qualiteg.
Loading an LLM from Hugging Face can take a very long time. In this article, we introduce a handy tool for exactly that situation.
When you download a Hugging Face LLM, you can see the progress, but once the model has been downloaded, you are left waiting anywhere from a few minutes to several tens of minutes for it to load. This is because it takes time to process the model data (weights and biases) from disk and load it into GPU VRAM. Have you ever found yourself anxiously wondering how far along that loading process actually is?
As a convenient feature, ChatStream can display the progress of this model loading time, as shown below.

The mechanism is quite simple: the processing time is measured on the first load, and when the same process is called a second time, a progress bar is displayed.
It is also easy to use. Simply wrap the model loading in LoadTime and the model will load with a progress bar.
Before
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16)
↓
↓
After
from chatstream import LoadTime
model = LoadTime(name=model_path,
fn=lambda: AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16))()
Full source code for loading the model
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) # the tokenizer is obtained after the model is loaded
Incidentally, this feature is also provided as a standalone library, so anyone can use it freely even without ChatStream.
Below is an introduction to the loadtime package.
How to Use loadtime
Installation
You can install LoadTime using pip.
pip install loadtime
Key Features
-
Real-time tracking: LoadTime provides real-time tracking of the loading process.
-
Progress bar: Displays a progress bar showing how much of the process has completed and how much remains.
-
Caching of past load times:
The time taken by the previous run is cached, and that cached information is used to provide the progress bar. -
Customizable display: LoadTime lets you customize the progress display with your own message.
Basic Usage
Here is some sample code.
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) # the tokenizer is obtained after the model is loaded
Initialization Parameters
| Parameter | Description |
|---|---|
| name | Specifies the name of the long-running process. When loading a Hugging Face model, specify the model name. |
| message | Specifies the message to display. If omitted, the default message is used. |
| pbar | When set to True, a progress bar and percentage are displayed. |
| dirname | Specifies the directory name where the cache is stored. |
| hf | When set to True, the timer display is used for loading Hugging Face models. If the model data has not yet been downloaded to disk, the Hugging Face loader displays the download progress, so this library does not display anything. |
| fn | Specifies the function that performs the long-running process. |
| fn_print | Specifies the function used for output. If omitted, output goes to the console. |