[QUESTION] Why we need modify the LinearWithGradAccumulationAndAsyncCommunication's backward function to support SP?

[b4b4o]

Thanks to authors for the great work!

As shown in this commit a84d634, It just seems to add a flag reshard_for_sequence_parallel for the Row Major Linear, The forward RS and backward AG of output_ = reduce_scatter_to_sequence_parallel_region(output_parallel) have been repackaged into backward().

I'm not sure why the change of this commit would support SP, or why it wouldn't support SP if without the change.

[ufotalent]

@b4b4o Hi thanks for the interest in our work! Without this change, the full size (after all-gether) grad_output tensor in SP is kept between B and release in W, which violates the optimization of SP to shard activation memory. After this change, we release the all-gathered grad_output after executing B, and re-all-gather the grad_output before executing W.