

TITLE
Training hangs at the very start while using deepspeed

COMMENTS
18

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.4.0
- base docker image: nvidia/cuda:10.1-cudnn7-devel-ubuntu18.04
- Python version: 3.8.8
- PyTorch version (GPU?): 1.7.1 (True)
- Tensorflow version (GPU?): 2.2.1 (True)
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: Yes, using deepspeed

### Who can help
@stas00 for deepspeed 

## Information

Model I am using Layoutlm:

I need to test my layoutlm model by training it only 1 epoch due to test purposes. However, training hangs at the very start without logging anything or returning an error message. When I disable deepspeed and launch my training with `python -m torch.distributed.launch` instead of `deepspeed --num_gpus={torch.cuda.device_count()} --num_nodes=1`, I manage to train for 1 epoch.

The tasks I am working on is:
* Token Classification

## To reproduce

I think it is a general issue. So, training any model with deepspeed for only one epoch may result in hanging process.


## Expected behavior

It would be possible to train a model only for 1 epoch not to waste time while testing. 

