Nonetype when using deepspeed 
## 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


Hi! I think the error you're looking for is actually a bit above to the error you're mentioning, namely:
```
RuntimeError: [enforce fail at CPUAllocator.cpp:67] . DefaultCPUAllocator: can't allocate memory: you tried to allocate 2329531392 bytes. Error code 12 (Cannot allocate memory)
```
which would indicate a memory issue!<|||||>What @LysandreJik said - check if you're close to using up all your gpu memory?

You can try to reduce `allgather_bucket_size` and `reduce_bucket_size` sizes,
https://huggingface.co/transformers/main_classes/trainer.html#zero
and see if it solves the problem.

In general if you see the traceback happening inside deepspeed most likely you will want to file an Issue at https://github.com/microsoft/DeepSpeed/ - Deepspeed is pretty much an independent engine.<|||||>This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.