[BUG] Can't perform gradient accumulation on pipeline parallelism with inputs of different lengths

Describe the bug

When performing pipeline parallelism, if gradient_accumulation_steps is >1, the engine will expect all inputs to be of the same shape. If we take an audio/text input as B x L x E I understand why the B batch dimension must be the same (for correct gradient accumulation across micro-batches?), but I see no reason why the length dimension L should be the same.

In the case of curriculum learning, this is only possible if dimensions differ only across batches, not across samples of a batch. And this is possible by calling reset_activation_shape() which can be run between batches only (after calling branch_train).

The issue in the code is in deepspeed/runtime/pipe/engine.py, where we have:

        if self.first_output_send:
            self.first_output_send = False
            self._send_tensor_meta(outputs, self.next_stage)

and

        if self.pipe_recv_buf is None:
            self.pipe_recv_buf = self._recv_tensor_meta(self.prev_stage)

This code only communicates the input dimensions/metadata for the first sample of the batch. Does not run when samples of the same batch differ in size. If a sample differs in size, within the same batch, the loss function will receive and input and a label of different sizes and will fail.

Is it possible to add a feature that allows for a similar call to reset_activation_shape between micro-batches?

Start in deepspeed/runtime/pipe/engine.py and trace the first_output_send and pipe_recv_buf paths, including _send_tensor_meta, _recv_tensor_meta, and reset_activation_shape(). Reproduce gradient accumulation with pipeline parallelism and inputs whose lengths differ within a batch. Done means the loss receives matching inputs and labels without requiring identical lengths across micro-batches.