yea, llama.cpp is it for the llama-style models.
About activation checkpointing — the way the LLM people do it is to intentionally sidestep autodiff and hand-write a clever backward. Sometimes people even do mathematically approximate stuff.

→ View original post on X — @soumithchintala