Selecting Which layers to recompute activations in?

[eliird]

Is there a way to control which layers to checkpoint when using recompute-granulariy?

If using transformer engine and flash attention, recompute-activation does not save any memory at all. The reason for that most likely is because flash attention is laready doing the recomputations for backward pass. (It is in the flash attention paper).

The above image explains how a bit more of memory can be saved, by checkpointing the input to FFN layer instead of the input to the flash attention layer. The idea in the above image is from the following paper. (https://arxiv.org/pdf/2310.03294v2)