[ChatStream] Adding a Progress Bar to Slow Model Loading

[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.

Read more