

TITLE
Running gpt-neo 2.7B with less than 13GB of system memory like Colab

COMMENTS
0

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 4
rocket: 0
eyes: 0

BODY
# 🚀 Feature request

A way to conserve regular system memory while loading large models. On systems without much system memory, the process crashes because it tries to load both the weight checkpoint and the model into system memory. In the case of gpt-neo 2.7B, this is even worse because the checkpoint is offered only in float32 format, taking up twice as much space.

## Motivation

Free Colab often offers GPUs with 16GB of VRAM but only about 13GB of RAM. This increases the barrier of entry for people to play with these kinds of models. I have found a way to load models in this kind of situation, but it is not general or well integrated. By posting it here, I hope that a more general implementation can be built at some point.

This would also help #11271.

## Your contribution

It is possible to work around this by loading the checkpoint directly into VRAM, casting it to float16, instantiating the model in VRAM and only then applying the weights from the checkpoint.

To do this, first a patch has to be applied to src/transformers/models/gpt_neo/modeling_gpt_neo.py. This is based on the 4.5.1 release.

    703c703
    <         self.h = nn.ModuleList([GPTNeoBlock(config, layer_id=i) for i in range(config.num_layers)])
    ---
    >         self.h = nn.ModuleList([GPTNeoBlock(config, layer_id=i).half().cuda() for i in range(config.num_layers)])
    890,891c890,891
    <         self.transformer = GPTNeoModel(config)
    <         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
    ---
    >         self.transformer = GPTNeoModel(config).half().cuda()
    >         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False).half().cuda()

This causes space for the biggest part of the model to be allocated directly on the GPU, which has more space in the free Colab scenario. It also moves the other parts of the model to GPU. Now the model can be instantiated like this:

    from transformers.file_utils import cached_path, WEIGHTS_NAME, hf_bucket_url
    model_name = "EleutherAI/gpt-neo-2.7B"
    archive_file = hf_bucket_url(model_name, filename=WEIGHTS_NAME)
    resolved_archive_file = cached_path(archive_file)
    checkpoint = torch.load(resolved_archive_file, map_location="cuda:0")
    for k in checkpoint.keys():
      checkpoint[k] = checkpoint[k].half()
    model = GPTNeoForCausalLM.from_pretrained(model_name, state_dict=checkpoint).half().to("cuda")
    for k in list(checkpoint.keys()):
      del checkpoint[k]