DeepSpeed

[zero-3] Handle forward parameter return correctly in nested cases

#2642

Merged