How Accelerate Loads Large Models

The traditional PyTorch model loading approach goes through a process of creating the model and then loading the weights into memory, which has the limitation of requiring CPU RAM far larger than the model size. For example, loading a 176B parameter model requires about 1.4TB of RAM.

Accelerate uses the following optimized process to solve this:

• First create an empty model without weights.
• Determine the device on which each layer will be placed.
• Repeat the process of partially loading weights, injecting them into the model, and moving them to the device.

The core of this technique lies in leveraging the meta device from PyTorch 1.9. Using the meta device, a model can be created with only the shape information of tensors and no actual data, allowing a huge model structure to be defined instantly without running into memory shortage issues.

Hugging Face supports this through the init_empty_weights context manager, enabling users to easily create empty models and efficiently load large models without modifying existing model code.