

TITLE
Nonetype when using deepspeed 

COMMENTS
3

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

BODY
## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.3.3
- Platform: linux
- Python version: 3.7
- PyTorch version (GPU?): 1.8
- Tensorflow version (GPU?): - 
- Using GPU in script?: - 
- Using distributed or parallel set-up in script?: - 

### Who can help
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @jplu

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

HF projects:

- nlp datasets: [different repo](https://github.com/huggingface/nlp)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->

deepspeed: @stas00

## Information

Hi there,
I am training run_mlm.py with wikipedia datasets with deepspeed, and for some datasets, I am getting this error below:
Do you have an idea why this might happen with deepspeed? looks like not to be a memory issue rather a None bug occuring:
```
  File "run_mlm.py", line 525, in <module>
    main()
  File "run_mlm.py", line 491, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/user/diba/libs/anaconda3/envs/transformer/lib/python3.7/site-packages/transformers/trainer.py", line 968, in train
    self.deepspeed.step()
  File "/user/diba/libs/anaconda3/envs/transformer/lib/python3.7/site-packages/deepspeed/runtime/engine.py", line 959, in step
    self._take_model_step(lr_kwargs)
  File "/user/diba/libs/anaconda3/envs/transformer/lib/python3.7/site-packages/deepspeed/runtime/engine.py", line 914, in _take_model_step
    self.optimizer.step()
  File "/user/diba/libs/anaconda3/envs/transformer/lib/python3.7/site-packages/deepspeed/runtime/zero/stage2.py", line 1425, in step
    self.optimizer.step()
  File "/user/diba/libs/anaconda3/envs/transformer/lib/python3.7/site-packages/torch/optim/optimizer.py", line 89, in wrapper
    return func(*args, **kwargs)
  File "/user/diba/libs/anaconda3/envs/transformer/lib/python3.7/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/user/diba/libs/anaconda3/envs/transformer/lib/python3.7/site-packages/torch/optim/adamw.py", line 121, in step
    group['eps'])
  File "/user/diba/libs/anaconda3/envs/transformer/lib/python3.7/site-packages/torch/optim/_functional.py", line 136, in adamw
    denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
RuntimeError: [enforce fail at CPUAllocator.cpp:67] . DefaultCPUAllocator: can't allocate memory: you tried to allocate 2329531392 bytes. Error code 12 (Cannot allocate memory)
Exception ignored in: <function tqdm.__del__ at 0x7f9b52ef4440>
Traceback (most recent call last):
  File "/user/diba/libs/anaconda3/envs/transformer/lib/python3.7/site-packages/tqdm/std.py", line 1090, in __del__
  File "/user/diba/libs/anaconda3/envs/transformer/lib/python3.7/site-packages/tqdm/std.py", line 1280, in close
  File "/user/diba/libs/anaconda3/envs/transformer/lib/python3.7/site-packages/tqdm/std.py", line 574, in _decr_instances
  File "/user/diba/libs/anaconda3/envs/transformer/lib/python3.7/site-packages/tqdm/_monitor.py", line 51, in exit
  File "/user/diba/libs/anaconda3/envs/transformer/lib/python3.7/threading.py", line 522, in set
  File "/user/diba/libs/anaconda3/envs/transformer/lib/python3.7/threading.py", line 365, in notify_all
  File "/user/diba/libs/anaconda3/envs/transformer/lib/python3.7/threading.py", line 348, in notify
TypeError: 'NoneType' object is not callable
```


## To reproduce
Sorry a bit hard to reproduce, I have modified a bit data collator of language modeling to let it do the T5 pretraining, which is currently not added in huggingface repo. If you could give me some advice from the error, it is greatly appreciated

