don't auto-recompute attention or linear

[t-vi]

Fixes: #1646

Thank you @kshitij12345 for the detailed issue.

[kshitij12345]

Overall looks good, just have one question. Also, we should add a test with a simple 2 layer model to verify sdpa or linear is not recomputed in backward.

Thank you @t-vi

[kshitij12345]

LGTM, thank you @t-vi

[IvanYashchuk]

That's a good quick fix, but can we revert to the previous default behavior: don't do any recomputation except for fused operations? Current logic doesn't use the information on whether the operation will be fused. To overwrite the default rule we can propagate torch.utils.checkpoint tags or have opt-in automatic passes.

[t-vi]

That's a good quick fix, but can we revert to the previous default behavior: don't do any recomputation except for fused operations? Current logic doesn't use the information on whether the operation will be fused. To overwrite the default rule we can propagate torch.utils.checkpoint tags or have opt-in automatic passes.

Maybe, I had that as one of the options in the issue, we went for this for now.

To my mind there are multiple parts:

• we want to be sure that we don't cause memory regressions,
• I think the rematerialization for forward and backward eventually needs to work without creating the joint trace.