Loading big models into memory

When loading a pre-trained model in PyTorch, the usual workflow looks like this:

import torch

my_model = ModelClass(...)
state_dict = torch.load(checkpoint_file)
my_model.load_state_dict(state_dict)

In plain English, those steps are:

1. Create the model with randomly initialized weights
2. Load the model weights (in a dictionary usually called a state dict) from the disk
3. Load those weights inside the model

While this works very well for regularly sized models, this workflow has some clear limitations when we deal with a huge model: in step 1, we load a full version of the model in RAM, and spend some time randomly initializing the weights (which will be discarded in step 3). In step 2, we load another full version of the model in RAM, with the pre-trained weights. If you’re loading a model with 6 billion parameters, this means you will need 24GB of RAM for each copy of the model, so 48GB in total (half of it to load the model in FP16).

How the Process Works: A Quick Overview

How the Process Works: Working with Code

Instantiating an empty model

The first tool Accelerate introduces to help with big models is a context manager init_empty_weights() that helps you initialize a model without using any RAM so that step 1 can be done on models of any size. Here is how it works:

from accelerate import init_empty_weights

with init_empty_weights():
    my_model = ModelClass(...)

For instance:

with init_empty_weights():
    model = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])

initializes an empty model with a bit more than 100B parameters. Behind the scenes, this relies on the meta device introduced in PyTorch 1.9. During the initialization under the context manager, each time a parameter is created, it is instantly moved to that device.

Sharded checkpoints

It’s possible your model is so big that even a single copy won’t fit in RAM. That doesn’t mean it can’t be loaded: if you have one or several GPUs, this is more memory available to store your model. In this case, it’s better if your checkpoint is split into several smaller files that we call checkpoint shards.

Accelerate will handle sharded checkpoints as long as you follow the following format: your checkpoint should be in a folder, with several files containing the partial state dicts, and there should be an index in the JSON format that contains a dictionary mapping parameter names to the file containing their weights. You can easily shard your model with save_model(). For instance, we could have a folder containing:

first_state_dict.bin
index.json
second_state_dict.bin

with index.json being the following file:

{
  "linear1.weight": "first_state_dict.bin",
  "linear1.bias": "first_state_dict.bin",
  "linear2.weight": "second_state_dict.bin",
  "linear2.bias": "second_state_dict.bin"
}

and first_state_dict.bin containing the weights for "linear1.weight" and "linear1.bias", second_state_dict.bin the ones for "linear2.weight" and "linear2.bias"

Loading weights

The second tool Accelerate introduces is a function load_checkpoint_and_dispatch(), that will allow you to load a checkpoint inside your empty model. This supports full checkpoints (a single file containing the whole state dict) as well as sharded checkpoints. It will also automatically dispatch those weights across the devices you have available (GPUs, CPU RAM), so if you are loading a sharded checkpoint, the maximum RAM usage will be the size of the biggest shard.

If you want to use big model inference with Transformers models, check out this documentation.

Here is how we can use this to load the GPT2-1.5B model.

Let’s download the sharded version of this model.

pip install huggingface_hub

from huggingface_hub import snapshot_download
checkpoint = "marcsun13/gpt2-xl-linear-sharded"
weights_location = snapshot_download(repo_id=checkpoint)

In order to initialize the model, we will use the library minGPT.

git clone https://github.com/karpathy/minGPT.git
pip install minGPT/

from accelerate import init_empty_weights
from mingpt.model import GPT

model_config = GPT.get_default_config()
model_config.model_type = 'gpt2-xl'
model_config.vocab_size = 50257
model_config.block_size = 1024

with init_empty_weights():
    model = GPT(model_config)

Then, load the checkpoint we just downloaded with:

from accelerate import load_checkpoint_and_dispatch

model = load_checkpoint_and_dispatch(
    model, checkpoint=weights_location, device_map="auto", no_split_module_classes=['Block']
)

By passing device_map="auto", we tell Accelerate to determine automatically where to put each layer of the model depending on the available resources:

• first, we use the maximum space available on the GPU(s)
• if we still need space, we store the remaining weights on the CPU
• if there is not enough RAM, we store the remaining weights on the hard drive as memory-mapped tensors

no_split_module_classes

This parameter will indicate that some of the modules with the name "Block" should not be split across different devices. You should set here all blocks that include a residual connection of some kind.

The device_map

You can see the device_map that Accelerate picked by accessing the hf_device_map attribute of your model:

model.hf_device_map

{'transformer.wte': 0,
 'transformer.wpe': 0,
 'transformer.drop': 0,
 'transformer.h.0': 0,
 ...
 'transformer.h.21': 0, 
 'transformer.h.22': 1, 
 'transformer.h.23': 1, 
 'transformer.h.24': 1,
 ...
 'transformer.h.47': 1, 
 'transformer.ln_f': 1, 
 'lm_head': 1}

It’s fully possible to create your own device map for the layers to use as well, specifying the GPU device to use (a number), "cpu", or "disk" and pass this in:

device_map = {
    "transformer.wte": "cpu",
    "transformer.wpe": 0,
    "transformer.drop": "cpu",
    "transformer.h.0": "disk"
}

model = load_checkpoint_and_dispatch(
    model, checkpoint=weights_location, device_map=device_map
)

Run the model

Now that we have done this, our model lies across several devices, and maybe the hard drive. But it can still be used as a regular PyTorch model:

from mingpt.bpe import BPETokenizer
tokenizer = BPETokenizer()
inputs = tokenizer("Hello, my name is").to(0)

outputs = model.generate(x1, max_new_tokens=10, do_sample=False)[0]
tokenizer.decode(outputs.cpu().squeeze())

Behind the scenes, Accelerate added hooks to the model, so that:

• at each layer, the inputs are put on the right device (so even if your model is spread across several GPUs, it works)
• for the weights offloaded on the CPU, they are put on a GPU just before the forward pass and cleaned up just after
• for the weights offloaded on the hard drive, they are loaded in RAM then put on a GPU just before the forward pass and cleaned up just after

This way, your model can run for inference even if it doesn’t fit on one of the GPUs or the CPU RAM!

Designing a device map

You can let Accelerate handle the device map computation by setting device_map to one of the supported options ("auto", "balanced", "balanced_low_0", "sequential") or create one yourself if you want more control over where each layer should go.

All the options will produce the same result when you don’t have enough GPU memory to accommodate the whole model (which is to fit everything that can on the GPU, then offload weights on the CPU or even on the disk if there is not enough RAM).

When you have more GPU memory available than the model size, here is the difference between each option:

• "auto" and "balanced" evenly split the model on all available GPUs, making it possible for you to use a batch size greater than 1.
• "balanced_low_0" evenly splits the model on all GPUs except the first one, and only puts on GPU 0 what does not fit on the others. This option is great when you need to use GPU 0 for some processing of the outputs, like when using the generate function for Transformers models
• "sequential" will fit what it can on GPU 0, then move on GPU 1 and so forth (so won’t use the last GPUs if it doesn’t need to).

First note that you can limit the memory used on each GPU by using the max_memory argument (available in infer_auto_device_map() and in all functions using it). When setting max_memory, you should pass along a dictionary containing the GPU identifiers (for instance 0, 1 etc.) and the "cpu" key for the maximum RAM you want to use for CPU offload. The values can either be an integer (in bytes) or a string representing a number with its unit, such as "10GiB" or "10GB".

Here is an example where we don’t want to use more than 10GiB on each of the two GPUs and no more than 30GiB of CPU RAM for the model weights:

from accelerate import infer_auto_device_map

device_map = infer_auto_device_map(my_model, max_memory={0: "10GiB", 1: "10GiB", "cpu": "30GiB"})

Additionally, if you do some additional operations with your outputs without placing them back on the CPU (for instance inside the generate method of Transformers) and if you placed your inputs on a GPU, that GPU will consume more memory than the others (Accelerate always place the output back to the device of the input). Therefore if you would like to optimize the maximum batch size and you have many GPUs, give the first GPU less memory. For example, with BLOOM-176B on 8x80 A100 setup, the close-to-ideal map is:

max_memory = {0: "30GIB", 1: "46GIB", 2: "46GIB", 3: "46GIB", 4: "46GIB", 5: "46GIB", 6: "46GIB", 7: "46GIB"}

as you can see we gave the remaining 7 GPUs ~50% more memory than GPU 0.

If you opt to fully design the device_map yourself, it should be a dictionary with keys being module names of your model and values being a valid device identifier (for instance an integer for the GPUs) or "cpu" for CPU offload, "disk" for disk offload. The keys need to cover the whole model, you can then define your device map as you wish: for instance, if your model has two blocks (let’s say block1 and block2) which each contain three linear layers (let’s say linear1, linear2 and linear3), a valid device map can be:

device_map = {"block1": 0, "block2": 1}

another one that is valid could be:

device_map = {"block1": 0, "block2.linear1": 0, "block2.linear2": 1, "block2.linear3": 1}

On the other hand, this one is not valid as it does not cover every parameter of the model:

device_map = {"block1": 0, "block2.linear1": 1, "block2.linear2": 1}

CPU offload only

If you want to offload your model on CPU, you can use cpu_offload(). As a result, all parameters of the model will be offloaded and only one copy of the state dict of the model will be kept. During the forward pass, parameters will be extracted from that state dict and put on the execution device and passed as they are needed, then offloaded again.

cpu_offload(model, execution_device)

You can also use cpu_offload_with_hook(). This function will offloads a model on the CPU and puts it back to an execution device when executed. The difference with cpu_offload() is that the model stays on the execution device after the forward and is only offloaded again when the offload method of the returned hook is called. Furthermore, cpu_offload_with_hook() is more performant but less memory saving. It is useful for pipelines running a model in a loop:

model_1, hook_1 = cpu_offload_with_hook(model_1, execution_device)
model_2, hook_2 = cpu_offload_with_hook(model_2, execution_device, prev_module_hook=hook_1)
model_3, hook_3 = cpu_offload_with_hook(model_3, execution_device, prev_module_hook=hook_2)

hid_1 = model_1(input)
for i in range(50):
    # model1 is offloaded on the CPU at the first iteration, model 2 stays on the GPU for this whole loop.
    hid_2 = model_2(hid_1)
# model2 is offloaded to the CPU just before this forward.
hid_3 = model_3(hid_3)

# For model3, you need to manually call the hook offload method.
hook_3.offload()

Disk offload only

To perform disk offload, you can use disk_offload(). As a result, all parameters of the model will be offloaded as memory-mapped array in a given folder. During the forward pass, parameters will be accessed from that folder and put on the execution device passed as they are needed, then offloaded again.

disk_offload(model, offload_dir, execution_device)

Limits and further development

We are aware of the current limitations in the API:

• infer_auto_device_map() (or device_map="auto" in load_checkpoint_and_dispatch()) tries to maximize GPU and CPU RAM it sees available when you execute it. While PyTorch is very good at managing GPU RAM efficiently (and giving it back when not needed), it’s not entirely true with Python and CPU RAM. Therefore, an automatically computed device map might be too intense on the CPU. Move a few modules to the disk device if you get crashes due to a lack of RAM.
• infer_auto_device_map() (or device_map="auto" in load_checkpoint_and_dispatch()) attributes devices sequentially (to avoid moving things back and forth) so if your first layer is bigger than the size of the GPU you have, it will end up with everything on the CPU/Disk.
• load_checkpoint_and_dispatch() and load_checkpoint_in_model() do not perform any check on the correctness of your state dict compared to your model at the moment (this will be fixed in a future version), so you may get some weird errors if trying to load a checkpoint with mismatched or missing keys.
• The model parallelism used when your model is split on several GPUs is naive and not optimized, meaning that only one GPU works at a given time and the other sits idle.
• When weights are offloaded on the CPU/hard drive, there is no pre-fetching (yet, we will work on this for future versions) which means the weights are put on the GPU when they are needed and not before.
• Hard-drive offloading might be very slow if the hardware you run on does not have fast communication between disk and CPU (like NVMes).