[BUG] PipelineEngine does not store losses on last stage in pipeline buffers

In the PipelineEngine class, the _exec_forward_pass function will store the loss (for the last pipeline stage) in a single variable self.loss for all micro-batches:
https://github.com/microsoft/DeepSpeed/blob/d24629f4fdaaa92df068de24f926d341f129112c/deepspeed/runtime/pipe/engine.py#L657-L663

This seems to be because the PipelineEngine assumes that self.loss will be differentiated immediately before any other forward passes are executed (as in 1F1B).
https://github.com/microsoft/DeepSpeed/blob/d24629f4fdaaa92df068de24f926d341f129112c/deepspeed/runtime/pipe/engine.py#L688-L691

This is problematic for implementing other pipeline schedules that do not behave this way (such as GPipe). One way to fix this would be to store the losses in the pipeline buffers and call backward on the correct loss (indexed by buffer_id).
https://github.com/microsoft/DeepSpeed/blob/d24629f4fdaaa92df068de24f926d341f129112c/deepspeed/runtime/pipe/engine.py#L150-L156

Start in deepspeed/runtime/pipe/engine.py at _exec_forward_pass, the backward path, and the pipeline buffer setup around the referenced lines. Trace how buffer_id maps to micro-batches and verify that each loss is retained and the matching loss is differentiated; done means schedules that delay backward, such as GPipe, no longer overwrite or use the wrong loss.