[QUESTION] Is there a way for Megatron to recompute the whole transformer layer except for the flash-attn part?

[xUhEngwAng]

Hi, as stated in the title, I'm wondering whether megatron provides native functionality like context_fn in torch.utils.checkpoint.checkpoint, such that flash-attn computation can be excluded from the recomputation of a transformer layer.

For now, I manully modified tensor_parallel.checkpoint to accept such an argument. However, issues remain when I want to capture the saved activations in flash-attn and offload them to cpu.

[sbhavani]

We are adding better CPU offloading infrastructure in Transformer Engine (related PR).

It doesn't directly solve your problem of selective recomputation but it provides significant improvements to activation offloading.

[xUhEngwAng]

We are adding better CPU offloading infrastructure in Transformer Engine (related PR).

It doesn't directly solve your problem of selective recomputation but it provides significant improvements to activation offloading.

Thanks for the quick reply! I tried the updated CPU offloading utility recently and it worked effectively. However, when I further tried to exclude the recomputation of flash-attn through context_fn, the added activation of flash-attn cannot be offloaded. Is this the expected bahaviour or is there something wrong the way I used it?

BTW, I use activation offloading along with full recomputation, as the seqlens in my case is ultra long. And I did observe the output of each transformer layer offloaded correctly according to pytorch memory profiler.