Megatron-DeepSpeed

4f42e591 - fix: grad_outputs instead of inputs, iterate over tuple