Recompute mlp

[ngc92]

Add the option to recompute the matmul in the MLP part.
Since this output is the largest remaining single tensor, we do get substantial improvements in memory consumption.

With 350M model at batch size 20, activation memory goes from 12042 MiB to 8362 MiB; conversely, on 2x4060Ti, I can increase the batch size from 20 to 28 without OOM. Overall, there is still a slowdown 30ktok/s to 28ktok/s, but
a) we cannot increase gradient accumulation arbitrarily due to bf16 rounding problems, so getting more tokens per fwd/bwd allows larger effective batch sizes
b) for smaller cards/larger models/longer contexts, this might make the difference between being able to run or not

so I think despite the noticeable drop in tok/s, this is an option that we want to have.