[PP] Why always cache chunk outputs for last stage ?

[mingyuanw-mt]

In the implementation of pipeline parallelism (take 1F1B as an example), the device correspond to the last stage will cache the output of forward into self.output_chunks for final output merge or reduction, which results the forward output of each micro batch size keeps alive during the step period and possibly retaining more GPU memory.

My question is Do we really always need to cache the forward outputs during training (If not, can PP make the caching behavior optional?)

I have posted a more detailed discuss on question-about-gpu-memory-usage-when-using-pipeline-parallelism-training-under-larger-micro-batch-count

Hi, @H-Huang @kwen2501 could you please share any suggestions you might have ？

cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta

[H-Huang]

This is good feedback and a feature we could implement. I am still curious about how it would affect overall max memory usage required (for the middle stages) if you could post some experimentation here that would be useful! (https://discuss.pytorch.org/t/question-about-gpu-memory-usage-when-using-pipeline-parallelism-training-under-larger-micro-batch-count/221886/2?u=h-huang

[wconstab]

I think we should implement this. Thoughts on the best API? should we add a 'return_outputs=False' flag to step API? (should it default to true so it doesn't break existing uses..?)

[mingyuanw-mt]

Hi, are there any updates on this? If not, I'd be happy to implement it.

[H-Huang]

Hi @mingyuanw-mt, im getting around to this. Let me send out a PR for this soon within the next few days.