[DeepSpeed] Features to integrate / Optimizations to add / Experiments to do

# 🚀 Feature request

While we have the support for main DeepSpeed features integrated, there are other powerful features that haven't been explored yet and which can provide even more various performance boosts. Some will probably require no changes on our side, while others require changes in the model and/or trainer.

This issue is to track what's possible and the priorities if any.

## Features to integrate

* [ ] [1-bit Adam](https://www.deepspeed.ai/tutorials/onebit-adam/) - Up to 5x less communication volume and up to 2x faster training
* [ ] [Progressive Layer Dropping](https://www.deepspeed.ai/tutorials/progressive_layer_dropping/) - Accelerating Training of Transformer-Based Language Models
* [ ] [DeepSpeed Sparse Attention](https://www.deepspeed.ai/tutorials/sparse-attention/) (Seems to be limited only to NVIDIA V100 )
* [ ] [DeepSpeed Transformer Kernel](https://www.deepspeed.ai/tutorials/transformer_kernel/) [api](https://deepspeed.readthedocs.io/en/latest/kernel.html)

Irrelevant to `transformers`:

* [ ] [DeepSpeed Activation Checkpointing](https://www.deepspeed.ai/docs/config-json/#activation-checkpointing) and extra discussion [here](https://github.com/microsoft/DeepSpeed/issues/665#issuecomment-760512582) - reduce the activation memory during model parallel training by partitioning activation checkpoints across model parallel GPUs, or offloading them to CPU. Since we don't use DS's PP there is no use for it.

## Experiments

Things to experiment with as well:

* [ ] try to profile model performance with DeepSpeed's `FlopsProfiler`

## Optimizations

* [ ] the new zero3 has a special requirement for inference with `--predict_with_generate` that all gpus run all `forward` calls even if they finished completing the predicted sequence early in `generate` - otherwise other gpus will hang waiting for the one that finished early. So currently the workaround is to simply always run till `max_length` in the `while` loop is reached. Which might be inefficient if we have a lot of short sequences, so need to use a synchronization trick to simultaneously quit the `while` loop when all gpus know it's safe to do so. @samyam posted a proof-of-concept for how to do that:

> We could maybe simplify by doing a single all_reduce, where gpus that are done will use a tensor with 0.0 and those that are not done will use 1.0. If the result of all reduce is 0.0 then everyone can stop, otherwise gpus that are done will do fake forward.
```
while sync.item() > 0.0:
p = model.forward(fake_input if am_i_done() else real_input)
sync =torch.tensor(0.0 if am_i_done() else 1.0)
torch.distributed.allreduce(sync)
```

At the moment this needs to be done in 5 places in the various search functions that `generate` may call.

For the full context please see: [this thread](https://github.com/microsoft/DeepSpeed/issues/860#issuecomment-799936583).

-------------------

If anybody would like to work on any of these items please open a dedicated issue so it'd be easier to track and please tag @stas00 to it.