[REQUEST] removing the requirement for all layers to always execute in sync

So by now we had to do multiple workaround to support ZeRO-3 in HF Transformers to deal with situations where one or more of the gpus wants to either finish early or randomly skip the layers. Here are the a few of these situations:

1. generate now includes a workaround that forces all gpus to continue generating even if one or more gpus finished early because their output was shorter than the other gpus:

https://github.com/huggingface/transformers/blob/69233cf03be5fbce0492f3997e139c4d05499e27/src/transformers/generation_utils.py#L1626-L1702

2. LayerDrop is another case now used in many models, where a layer is randomly dropped, and the workaround is to continue running all layers always but randomly ignoring the outcome:

https://github.com/huggingface/transformers/blob/69233cf03be5fbce0492f3997e139c4d05499e27/src/transformers/models/wav2vec2/modeling_wav2vec2.py#L817-L849

And there are other situation.

The workarounds are OK, except they waste resources, and make the code much more difficult to read and it's going in the opposite direction of ease of use, as the models no longer work out of box.

I was thinking if perhaps Deepspeed could provide a built-in solution for these situations. After all the only reason for forcing the forward call is to gather the weights in the pre hook. And if this doesn't happen Deepspeed usually either hangs or there is an error about the wrong sequence of parameters. So perhaps there could be a way to run the pre-hook but not the forward itself? (and backward where needed).

Perhaps in such cases instead of the normal forward it could call a special deepspeed forward which will just gather the weights and not need to run the normal forward. Except I think the problem emerges if that forward had a bunch of other forwards called from inside of it. And I suppose that due to the dynamic nature of most models, there is no way to "trace" and know ahead of time which path to "skip". Perhaps it's not quite possible. But I thought I'd ask your genius if perhaps you have some ideas to solve this transparently to the model.

e.g. one of the most recent reports with LayerDrop https://github.com/huggingface/transformers/issues/16688

Thank you!

@samyam, @jeffra

Start with the referenced Transformers paths, generation_utils.py and modeling_wav2vec2.py, to understand the early-finish and LayerDrop workarounds. Then inspect DeepSpeed's ZeRO-3 forward pre-hook behavior and parameter-order requirements. Done would mean a defined, transparent approach for divergent layer execution that avoids hangs and unnecessary computation in these cases.