[FSDP] fix for high GPU reserved memory (v2)

[ruanslv]

v2 for #972

The general idea is to try to make a guess on the FSDP module execution order and store that in two lists (one for the forward and another for the backward pass) during the first pass. Then, as opposed to letting the CPU run free and scheduling all GPU operations ahead of time (while reserving GPU memory for each), we only schedule the all-gather for the next module and wait until computations for the current module are finished before continuing. By scheduling the all-gather for the next module, we attempt to keep a good parallelism between data transfer and compute streams.

Given the lack of a global/static execution graph in Pytorch, my understanding on the best we can do here is a heuristic based on local information. Pytorch Distributed is also facing a similar issue and, given there is no perfect solution, arguments for/against each approach largely depend on the models you pick to measure success.

Advantages of current approach:
(1) Based on several comments in the original PR, it showed to significantly help different large scale runs that are memory bound. I'm attaching some profiles for a ~370M param transformer showing the new behaviour. As long as we get the execution order right and there is not a lot of variance in terms of how long each module takes to run, the parallelism across compute and data transfer is maintained.
(2) We only wait until computations are finished for the modules where reshard_after_forward = True, with the assumption that if this is not set then memory is not a limitation and we should let CPU continue. This should help preventing side-effects on models that do not care about memory.

Disadvantages:
(1) We may not always schedule the correct all-gather, which can cause run delays. One example is with activation checkpointing, where the execution needs to go further back in the model and execute a forward to recompute activations. The first module in this forward is not all-gathered ahead of time. Another example is if your execution order changes across different passes. So it is theoretically possible for this PR to cause performance degradations in some scenarios.
(2) We may under-schedule all-gathers, which can also cause performance degradation. This could happen on models where there is a big variation on execution time across different modules, and we end up having to wait for a long all-gather to finish after running a module that had very short computation time.

One alternative heuristic proposed by Min was to use the amount of available memory to make a call on whether to wait for the current module to finish or continue execution. However, I found it hard to find a justifiable memory threshold (either relative or absolute) that would work well across a large variety of cases, especially with a lack of a comprehensive benchmark to experiment with. Given we've already seen examples of the approach in this PR working well in practice, it seems safer to just go with this route instead and reevaluate if we find evidence of the contrary.

Original behavior: computation stream is always active, however CPU schedules everything at once.

New behavior: computation stream still mostly active, with CPU scheduling one module at a time.

[min-xu-ai]

@ruanslv, thanks for the detailed description and explanation in this new PR! Kudos for adding a new test.

I haven't looked at the code changes in details but can you say a bit about the changes in the overlap fsdp test? That test was put in to specifically catch a performance issue. With the changes you make in that test, does it still catch the performance issue it intended to catch? You can "git log" on the file and find out which issue was it introduced to address. There was a bug fix together with the test I believe. You can probably undo the fix and see if the test catch the issue or not.

[ruanslv]

Sorry had missed SSD offload tests as they only run on a later Pytorch version compared to the one I was using. Fixed them with the last commit. The remaining CI failures are unrelated.

[ruanslv]

@min-xu-ai re: the overlap test, in the deleted part it was measuring that CPU is running ahead of GPU, which is exactly what we are disabling in this PR so that behaviour is not expected anymore. Further down, the test still measures whether all-gather and compute are running in parallel, which still should happen here. I couldn't think of relevant tests comparing CPU against GPU timing, happy to take suggestions

Original commit for reference: 8a42a8e

[zhaojuanmao]

@ruanslv the traces are amazing and clearly showed the root cause is because of fast CPU thread. @mrshenli this also proved what we've discussed.

This seems to be a common problem for cudaCacheAllocator? it seems that FSDP use cases are hurt more by high reserved memory issue because these use cases are usually training large models, when reserved memory is high and fragmented, allocating a new big tensor from caches tends to fail more frequently and thus results in either OOM or slow cudaMalloc (a lot of cudaFree calls first).

@mrshenli will the pluggable cudaCacheAllocator help resolve the problem?

[awgu]

Hi @ruanslv! Do you have a code pointer/hardware setup info for me to reproduce the same 370M parameter transformer and check the trace (for a different FSDP patch)? Also, are there any publicly available setups for replicating the high GPU reserved memory?

I think I have a similar but slightly different approach for this.

[stephenroller]

Does it always have issues with activation checkpointing? Is there a way we might address those? Activation checkpointing is pretty much non-optional for any any model 5B+ params.

[ruanslv]

Sorry folks, I missed the comments here. Please feel free to ping me on work chat directly too.

@awgu I used an internal benchmark we have at https://github.com/fairinternal/fairscale_benchmarks/blob/main/benchmarks/transformer_wikitext2_benchmark.py. This model doesn't necessarily help seeing the issue, I used it as an illustration / test environment for prefetching and CPU behavior. I'm not aware of any public setups for replicating the problem.

@stephenroller I couldn't find a way with prefetching to get around checkpointing. However, @zhaojuanmao organized a chat where we concluded that the optimal solution here is not to use pretetch + synchronize(). Instead, we should use explicit CUDA events to synchronize behaviour across different CUDA streams and help the CUDA cache allocator behave properly (i.e. reusing memory as opposed to reserving new one based on the CUDA event information).

I haven't found the cycles to explore this idea yet, but if it works it would be a generic / clean solution that doesn't require CPU waits.