gpt-neo-1-3B large memory increase during training even with a small training dataset
Author: sandeeppagey
Pull request: no
Labels: name: bug

### System Info

```shell
- `transformers` version: 4.19.2
- Platform: Linux-5.4.0-1075-aws-x86_64-with-debian-buster-sid
- Python version: 3.7.10
- Huggingface_hub version: 0.5.1
- PyTorch version (GPU?): 1.8.1+cu111 (True)
- Tensorflow version (GPU?): 2.8.0 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: NO
- Using distributed or parallel set-up in script?: NO
```


### Who can help?

@patil-suraj 
Training the gpt-neo-1-3B model with a training dataset of 100 sentences (about 75 words max for each sentence, say about 100 tokens), there is jump of close to 29GB in memory consumption during the training. Tried batch sizes of 8/4/2 showing similar increase. Is that expected, or does this memory increase needs further investigation? Model is loaded without setting any of return_dict, output_attentions or output_hidden_states to True.

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Simple training loop as shown in code example in
https://towardsdatascience.com/how-to-fine-tune-gpt-2-for-text-generation-ae2ea53bc272

Use gpt-neo-1-3b and 100 entries for training.

Doing training for 1 epoch itself and the memory jumps by about 29G.

%memit model = train(dataset, model, tokenizer, epochs=1, gpt2_type="gpt_neo_1_3B")

peak memory: 36634.16 MiB, increment: 29428.80 MiB


### Expected behavior

```shell
Not sure if there is some issue with code or this kind of memory jump is expected.
```


Comments:
Hi @sandeeppagey !

Could you post a simple code-snippet here, with more details such as the GPU memory, the `dtype` for training. 

Also note that, the model weights in fp32 are about 5GB and Adam takes 3x more memory, so this alone is about ~20GB and then there's more memory to store the activations and gradients. 

Hi @patil-suraj 

Thanks for the explanation. I am not running on GPU. I am running
on CPU just to check the memory increase. I have not changed the dtype so I guess it is fp32.
I agree with your calculations that the model training will require about 20G of memory. Please close
the issue.
The relevant code snippet is as follows. 

model = GPTNeoForCausalLM.from_pretrained('/home1/pretrained_models/gpt_neo_1_3B')

    optimizer = AdamW(model.parameters(), lr=2e-5)
    scheduler = get_linear_schedule_with_warmup(
        optimizer, num_warmup_steps=200, num_training_steps=-1
    )

    train_dataloader = DataLoader(dataset, batch_size=1, shuffle=True)
    loss=0
    accumulating_batch_count = 0
    input_tensor = None

    for epoch in range(1):

        print(f"Training epoch {epoch}")
        print(f"Loss before: {loss}")
        for idx, entry in tqdm(enumerate(train_dataloader)):
            input_tensor = entry #tokenization has already been done

            outputs = model(input_tensor, labels=input_tensor)
            #print(isinstance(outputs, dict))
            #print(type(outputs))
            loss = outputs[0]
            loss.backward()
            
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()
            model.zero_grad()