Lazy initialization

Author: Hongxin Liu

Prerequisite:

• Train with booster

Introduction

Lazy initialization defers model initialization. It saves memory when initializing large models.

If your model has N billion parameters and your memory (or GPU memory) is M GB, we recommend you use lazy initialization when 4N >= M. Otherwise, it is optional.

Usage

Lazy initialization must be used with booster.

API reference

{{ autodoc:colossalai.lazy.LazyInitContext }}

Example

import colossalai
from colossalai.lazy import LazyInitContext
from colossalai.booster import Booster
from colossalai.booster.plugin import GeminiPlugin

from transformers import LlamaForCausalLM, LlamaConfig, BertForPreTraining

colossalai.launch()
plugin = GeminiPlugin()
booster = Booster(plugin)

# 1. Initialize model from scratch
# Initialization on cuda will accelerate the initialization process but take more GPU memory.
with LazyInitContext(default_device="cuda"):
    model = LlamaForCausalLM(LlamaConfig(hidden_size=64, intermediate_size=172, num_hidden_layers=4, num_attention_heads=4))
model, *_ = booster.boost(model)

# 2. Initialize model from pretrained
with LazyInitContext():
    model = BertForPreTraining.from_pretrained("prajjwal1/bert-tiny")
model, *_ = booster.boost(model)

⚠️ Lazy initialization from pretrained is supported for colossalai>0.3.3 or main branch.

Limitations

As we claimed, lazy initialization must be used with booster. And only several plugins support it.

Not all models can be lazily initialized. In some cases, a part of parameters/buffers may be early initialized. But don't worry, this part usually takes a small proportion of the whole model.

And some models are not supported at all which will raise an error. We tested models in torchvision, diffusers, timm, transformers, torchaudio and torchrec. Below models are not supported: